├── .github └── workflows │ ├── ci.yml │ └── publish.yml ├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── LICENSE ├── README.md ├── cli ├── Cargo.toml └── src │ ├── commands │ ├── build.rs │ ├── debug.rs │ ├── login.rs │ ├── mod.rs │ ├── predict.rs │ └── push.rs │ ├── config.rs │ ├── context.rs │ ├── docker │ ├── auth.rs │ ├── builder.rs │ ├── dockerfile.rs │ ├── mod.rs │ └── predictor.rs │ ├── helpers.rs │ ├── main.rs │ └── templates │ └── Dockerfile ├── core ├── Cargo.toml └── src │ ├── http.rs │ ├── lib.rs │ └── spec.rs ├── examples ├── blur │ ├── Cargo.lock │ ├── Cargo.toml │ └── src │ │ └── main.rs ├── hello-world │ ├── Cargo.lock │ ├── Cargo.toml │ └── src │ │ └── main.rs ├── resnet │ ├── .gitignore │ ├── Cargo.lock │ ├── Cargo.toml │ ├── README.md │ ├── src │ │ └── main.rs │ └── weights │ │ └── .gitkeep └── stable-diffusion │ ├── .gitignore │ ├── Cargo.lock │ ├── Cargo.toml │ ├── README.md │ ├── src │ └── main.rs │ └── weights │ └── .gitkeep ├── lib ├── Cargo.toml └── src │ ├── errors.rs │ ├── helpers │ ├── headers.rs │ ├── mod.rs │ └── openapi.rs │ ├── lib.rs │ ├── prediction.rs │ ├── routes │ ├── docs.rs │ ├── mod.rs │ ├── predict.rs │ └── system.rs │ ├── runner.rs │ ├── server.rs │ ├── shutdown.rs │ ├── spec.rs │ └── webhooks.rs └── rustfmt.toml /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | on: [push, pull_request] 2 | 3 | name: CI 4 | 5 | jobs: 6 | check: 7 | name: Lint 8 | runs-on: ubuntu-latest 9 | steps: 10 | - name: Checkout repository 11 | uses: actions/checkout@v3 12 | - uses: actions-rs/toolchain@v1 13 | with: 14 | profile: minimal 15 | toolchain: nightly 16 | override: true 17 | components: rustfmt, clippy 18 | - name: Cache build 19 | uses: Swatinem/rust-cache@v2 20 | with: 21 | key: cache 22 | - name: Check formatting 23 | uses: actions-rs/cargo@v1 24 | with: 25 | command: fmt 26 | args: --all -- --check 27 | - name: Clippy 28 | uses: actions-rs/clippy-check@v1 29 | with: 30 | token: ${{ secrets.GITHUB_TOKEN }} 31 | args: --all-features --all-targets 32 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to Cargo 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | 7 | jobs: 8 | publish: 9 | runs-on: ubuntu-latest 10 | environment: cargo 11 | 12 | env: 13 | CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} 14 | 15 | steps: 16 | - uses: actions/checkout@v3 17 | with: 18 | fetch-depth: 2 19 | 20 | - name: Cache build 21 | uses: Swatinem/rust-cache@v2 22 | with: 23 | key: cache 24 | 25 | - name: Publish Crate 26 | uses: seunlanlege/cargo-auto-publish@2 27 | with: 28 | toml: core/Cargo.toml lib/Cargo.toml cli/Cargo.toml 29 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = ["core", "lib", "cli"] 3 | exclude = ["examples"] 4 | resolver = "2" 5 | 6 | [workspace.package] 7 | edition = "2021" 8 | readme = "README.md" 9 | license = "Apache-2.0" 10 | keywords = ["cog", "machine-learning"] 11 | categories = ["science", "command-line-utilities"] 12 | repository = "https://github.com/m1guelpf/cog-rust" 13 | authors = ["Miguel Piedrafita "] 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Cog[-rust]: Containers for machine learning 2 | 3 | Cog is an open-source tool that lets you package Rust ML models in a standard, production-ready container. 4 | 5 | It's output should be interchangeable with [Replicate's own Cog](https://github.com/replicate/cog) (for Python models). 6 | 7 | ## Highlights 8 | 9 | - 📦 **Docker containers without the pain.** Writing your own `Dockerfile` can be a bewildering process. With Cog, you define your environment [inside your Cargo.toml](#how-it-works) and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, minimal image sizes, sensible environment variable defaults, and so on. 10 | 11 | - 🤬️ **No more CUDA hell.** Cog knows which CUDA/cuDNN/tch/tensorflow combos are compatible and will set it all up correctly for you. 12 | 13 | - ✅ **Define the inputs and outputs for your model in Rust.** Then, Cog generates an OpenAPI schema and validates the inputs and outputs with JSONSchema. 14 | 15 | - 🎁 **Automatic HTTP prediction server**: Your model's types are used to dynamically generate a RESTful HTTP API using [axum](https://github.com/tokio-rs/axum). 16 | 17 | - ☁️ **Cloud storage.** Files can be read and written directly to Amazon S3 and Google Cloud Storage. (Coming soon.) 18 | 19 | - 🚀 **Ready for production.** Deploy your model anywhere that Docker images run. Your own infrastructure, or [Replicate](https://replicate.com). 20 | 21 | ## How it works 22 | 23 | Easily define your environment inside your `Cargo.toml`. Cog infers the rest: 24 | 25 | ```toml 26 | [package] 27 | name = "ml-model" 28 | 29 | [package.metadata.cog] 30 | gpu = true # optional, defaults to false 31 | image = "docker-image-name" # optional, defaults to `cog-[package.name]` 32 | ``` 33 | 34 | Define how predictions are run on your model on your `main.rs`: 35 | 36 | ```rust 37 | use anyhow::Result; 38 | use cog_rust::Cog; 39 | use schemars::JsonSchema; 40 | use std::collections::HashMap; 41 | use tch::{ 42 | nn::{ModuleT, VarStore}, 43 | vision::{imagenet, resnet::resnet50}, 44 | Device, 45 | }; 46 | 47 | #[derive(serde::Deserialize, schemars::JsonSchema)] 48 | struct ModelRequest { 49 | /// Image to classify 50 | image: cog_rust::Path, 51 | } 52 | 53 | struct ResnetModel { 54 | model: Box, 55 | } 56 | 57 | impl Cog for ResnetModel { 58 | type Request = ModelRequest; 59 | type Response = HashMap; 60 | 61 | async fn setup() -> Result { 62 | let mut vs = VarStore::new(Device::Cpu); 63 | vs.load("weights/model.safetensors")?; 64 | let model = Box::new(resnet50(&vs.root(), imagenet::CLASS_COUNT)); 65 | 66 | Ok(Self { model }) 67 | } 68 | 69 | fn predict(&self, input: Self::Request) -> Result { 70 | let image = imagenet::load_image_and_resize224(&input.image)?; 71 | let output = self 72 | .model 73 | .forward_t(&image.unsqueeze(0), false) 74 | .softmax(-1, tch::Kind::Float); 75 | 76 | Ok(imagenet::top(&output, 5) 77 | .into_iter() 78 | .map(|(prob, class)| (class, 100.0 * prob)) 79 | .collect()) 80 | } 81 | } 82 | 83 | cog_rust::start!(ResnetModel); 84 | ``` 85 | 86 | Now, you can run predictions on this model: 87 | 88 | ```console 89 | $ cargo cog predict -i @input.jpg 90 | --> Building Docker image... 91 | --> Running Prediction... 92 | --> Output written to output.jpg 93 | ``` 94 | 95 | Or, build a Docker image for deployment: 96 | 97 | ```console 98 | $ cargo cog build -t my-colorization-model 99 | --> Building Docker image... 100 | --> Built my-colorization-model:latest 101 | 102 | $ docker run -d -p 5000:5000 --gpus all my-colorization-model 103 | 104 | $ curl http://localhost:5000/predictions -X POST \ 105 | -H 'Content-Type: application/json' \ 106 | -d '{"input": {"image": "https://.../input.jpg"}}' 107 | ``` 108 | 109 | ## Why am I building this? 110 | 111 | The Replicate team has done an amazing job building the simplest way to go from Python notebook to Docker image to API endpoint. 112 | 113 | However, using Python as the base layer comes with its on share of challenges, like enormus image sizes or extra latency on model requests. 114 | 115 | As the non-Python ML ecosystem slowly flourishes (see [whisper.cpp](https://github.com/ggerganov/whisper.cpp) and [llama.cpp](https://github.com/ggerganov/llama.cpp) for example), cog-rust will provide that extra performance exposed on the same interfaces users and tools are already used to. 116 | 117 | ## Prerequisites 118 | 119 | - **macOS, Linux or Windows**. Cog works anywhere Rust works. 120 | - **Docker**. Cog uses Docker to create a container for your model. You'll need to [install Docker](https://docs.docker.com/get-docker/) before you can run Cog. 121 | 122 | ## Install 123 | 124 | 125 | 126 | You can install Cog with Cargo: 127 | 128 | ```console 129 | cargo install cargo-cog 130 | ``` 131 | 132 | ## Usage 133 | 134 | ``` 135 | $ cargo cog --help 136 | A cargo subcommand to build, run and publish machine learning containers 137 | 138 | Usage: cargo cog [OPTIONS] [COMMAND] 139 | 140 | Commands: 141 | login Log in to Replicate's Docker registry 142 | build Build the model in the current directory into a Docker image 143 | push Build and push model in current directory to a Docker registry 144 | predict Run a prediction 145 | help Print this message or the help of the given subcommand(s) 146 | 147 | Options: 148 | -h, --help Print help 149 | -V, --version Print version 150 | ``` 151 | -------------------------------------------------------------------------------- /cli/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cargo-cog" 3 | version = "1.0.10" 4 | description = "A cargo subcommand to build, run and publish machine learning containers" 5 | readme = { workspace = true } 6 | edition = { workspace = true } 7 | authors = { workspace = true } 8 | license = { workspace = true } 9 | keywords = { workspace = true } 10 | categories = { workspace = true } 11 | repository = { workspace = true } 12 | 13 | [dependencies] 14 | dirs = "5.0.1" 15 | indoc = "2.0.1" 16 | serde = "1.0.164" 17 | anyhow = "1.0.71" 18 | base64 = "0.21.2" 19 | map-macro = "0.2.6" 20 | schemars = "0.8.12" 21 | thiserror = "1.0.40" 22 | mime_guess = "2.0.4" 23 | serde_json = "1.0.96" 24 | webbrowser = "0.8.10" 25 | cargo_metadata = "0.15.4" 26 | cog-core = { path = "../core", version = "0.2.0" } 27 | clap = { version = "4.3.3", features = ["derive"] } 28 | tokio = { version = "1.28.2", features = ["full"] } 29 | reqwest = { version = "0.11.18", features = ["json"] } 30 | dataurl = { version = "0.1.2", default-features = false } 31 | tree_magic_mini = { version = "3.0.3", features = ["with-gpl-data"] } 32 | -------------------------------------------------------------------------------- /cli/src/commands/build.rs: -------------------------------------------------------------------------------- 1 | use crate::Context; 2 | 3 | pub fn handle(ctx: Context, tag: Option) { 4 | let image_name = ctx.into_builder().build(tag); 5 | println!("Image built as {image_name}"); 6 | } 7 | -------------------------------------------------------------------------------- /cli/src/commands/debug.rs: -------------------------------------------------------------------------------- 1 | use crate::Context; 2 | 3 | pub fn handle(ctx: Context) { 4 | println!("{}", ctx.into_builder().generate_dockerfile()); 5 | } 6 | -------------------------------------------------------------------------------- /cli/src/commands/login.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use reqwest::StatusCode; 3 | use serde_json::Value; 4 | 5 | use crate::{ 6 | docker, 7 | helpers::{load_from_stdin, wait_for_input}, 8 | }; 9 | 10 | pub async fn handle(token_stdin: bool, registry: String) { 11 | let token = if token_stdin { 12 | load_from_stdin() 13 | } else { 14 | get_token_interactive(®istry).await 15 | }; 16 | 17 | let username = verify_token(®istry, &token).await.unwrap(); 18 | 19 | docker::store_credentials(®istry, &username, &token).unwrap(); 20 | println!("You've successfully authenticated as {username}! You can now use the '{registry}' registry."); 21 | } 22 | 23 | async fn get_token_interactive(registry: &str) -> String { 24 | let token_url = get_display_token_url(registry).await.unwrap(); 25 | 26 | println!("This command will authenticate Docker with Replicate's '{registry}' Docker registry. You will need a Replicate account."); 27 | println!("Hit enter to get started. A browser will open with an authentication token that you need to paste here."); 28 | let _ = wait_for_input(); 29 | println!("If it didn't open automatically, open this URL in a web browser:\n{token_url}"); 30 | let _ = webbrowser::open(&token_url); 31 | println!("Once you've signed in, copy the authentication token from that web page, paste it here, then hit enter:"); 32 | 33 | wait_for_input() 34 | } 35 | 36 | async fn get_display_token_url(registry_host: &str) -> Result { 37 | let resp = reqwest::get(format!( 38 | "{}/cog/v1/display-token-url", 39 | if registry_host.contains("://") { 40 | registry_host.to_string() 41 | } else { 42 | format!("https://{registry_host}") 43 | }, 44 | )) 45 | .await 46 | .unwrap(); 47 | 48 | if matches!(resp.status(), StatusCode::NOT_FOUND) { 49 | return Err(format!( 50 | "{registry_host} is not the Replicate registry\nPlease log in using 'docker login'", 51 | )); 52 | } 53 | 54 | if !resp.status().is_success() { 55 | return Err(format!( 56 | "{registry_host} returned HTTP status {}", 57 | resp.status() 58 | )); 59 | } 60 | 61 | let body: Value = resp.json().await.unwrap(); 62 | Ok(body 63 | .get("url") 64 | .and_then(|v| Some(v.as_str()?.to_string())) 65 | .unwrap()) 66 | } 67 | 68 | async fn verify_token(registry_host: &str, token: &str) -> Result { 69 | let resp = reqwest::Client::new() 70 | .post(format!( 71 | "{}/cog/v1/verify-token", 72 | if registry_host.contains("://") { 73 | registry_host.to_string() 74 | } else { 75 | format!("https://{registry_host}") 76 | }, 77 | )) 78 | .form(&[("token", token)]) 79 | .send() 80 | .await 81 | .unwrap(); 82 | 83 | if matches!(resp.status(), StatusCode::NOT_FOUND) { 84 | return Err("User does not exist".to_string()); 85 | } 86 | 87 | if !resp.status().is_success() { 88 | return Err(format!( 89 | "Failed to verify token, got status {}", 90 | resp.status() 91 | )); 92 | } 93 | 94 | let body: Value = resp.json().await.unwrap(); 95 | Ok(body 96 | .get("username") 97 | .and_then(|v| Some(v.as_str()?.to_string())) 98 | .unwrap()) 99 | } 100 | -------------------------------------------------------------------------------- /cli/src/commands/mod.rs: -------------------------------------------------------------------------------- 1 | use clap::Subcommand; 2 | use std::path::PathBuf; 3 | 4 | use crate::Context; 5 | 6 | mod build; 7 | mod debug; 8 | mod login; 9 | mod predict; 10 | mod push; 11 | 12 | #[derive(Debug, Subcommand)] 13 | #[clap(bin_name = "cargo-cog", version, author)] 14 | pub enum Command { 15 | /// Log in to Replicate's Docker registry 16 | Login { 17 | /// Pass login token on stdin instead of opening a browser. You can find your Replicate login token at https://replicate.com/auth/token 18 | #[clap(long)] 19 | token_stdin: bool, 20 | /// Registry host 21 | #[clap(hide = true, default_value = "r8.im")] 22 | registry: String, 23 | }, 24 | /// Generate a Dockerfile for your project 25 | #[clap(hide = true)] 26 | Debug, 27 | 28 | /// Build the model in the current directory into a Docker image 29 | Build { 30 | /// A name for the built image in the form 'repository:tag' 31 | #[clap(short, long)] 32 | tag: Option, 33 | }, 34 | 35 | /// Build and push model in current directory to a Docker registry 36 | Push { 37 | /// A name for the built image 38 | image: Option, 39 | }, 40 | 41 | /// Run a prediction. 42 | Predict { 43 | /// Run the prediction on this Docker image (it must be an image that has been built by Cog). 44 | /// 45 | /// Will build and run the model in the current directory if left empty 46 | image: Option, 47 | 48 | /// Output path 49 | #[clap(short)] 50 | output: Option, 51 | 52 | /// Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg 53 | #[clap(short)] 54 | input: Option>, 55 | }, 56 | } 57 | 58 | pub async fn exec(command: Command, ctx: Context) { 59 | match command { 60 | Command::Debug => debug::handle(ctx), 61 | Command::Build { tag } => build::handle(ctx, tag), 62 | Command::Push { image } => push::handle(ctx, &image), 63 | Command::Login { 64 | registry, 65 | token_stdin, 66 | } => login::handle(token_stdin, registry).await, 67 | Command::Predict { 68 | image, 69 | output, 70 | input, 71 | } => predict::handle(ctx, image, input, output).await, 72 | }; 73 | } 74 | -------------------------------------------------------------------------------- /cli/src/commands/predict.rs: -------------------------------------------------------------------------------- 1 | use base64::{engine::general_purpose::STANDARD as Base64, Engine}; 2 | use cog_core::http::Status; 3 | use dataurl::DataUrl; 4 | use mime_guess::Mime; 5 | use schemars::schema::SchemaObject; 6 | use serde_json::Value; 7 | use std::{collections::HashMap, fs, path::PathBuf, str::FromStr}; 8 | 9 | use crate::{ 10 | docker::{Docker, Predictor}, 11 | Context, 12 | }; 13 | 14 | pub async fn handle( 15 | ctx: Context, 16 | image: Option, 17 | inputs: Option>, 18 | output: Option, 19 | ) { 20 | let image = image.map_or_else( 21 | || ctx.clone().into_builder().build(None), 22 | |image| { 23 | if Docker::inspect_image(&image).is_err() { 24 | Docker::pull(&image).unwrap(); 25 | } 26 | 27 | image 28 | }, 29 | ); 30 | 31 | println!("Starting Docker image {image} and running setup()..."); 32 | let mut predictor = Predictor::new(image); 33 | 34 | predictor.start().await; 35 | predict_individual_inputs(&predictor, inputs, output).await; 36 | } 37 | 38 | async fn predict_individual_inputs( 39 | predictor: &Predictor, 40 | inputs: Option>, 41 | mut output: Option, 42 | ) { 43 | println!("Running prediction..."); 44 | let schema = predictor.get_schema().unwrap(); 45 | let inputs = inputs 46 | .map(|inputs| parse_inputs(inputs, &schema)) 47 | .unwrap_or_default(); 48 | 49 | let prediction = predictor.predict(inputs).await.unwrap(); 50 | 51 | if !matches!(prediction.status, Status::Succeeded) { 52 | eprintln!( 53 | "Prediction failed: {}", 54 | prediction.error.unwrap_or_default() 55 | ); 56 | std::process::exit(1); 57 | } 58 | 59 | let response_schema = (|| { 60 | schema 61 | .extensions 62 | .get("components")? 63 | .get("schemas")? 64 | .get("Output") 65 | })() 66 | .unwrap(); 67 | 68 | let out = parse_response(&prediction.output.unwrap(), response_schema, &mut output); 69 | 70 | match out { 71 | SerializedResponse::Text(text) => println!("{text}"), 72 | SerializedResponse::Bytes(bytes) => { 73 | let output = output.expect("No output file specified"); 74 | 75 | fs::write(&output, bytes).expect("Failed to write output file"); 76 | println!("Written output to {}", output.display()); 77 | }, 78 | } 79 | } 80 | 81 | #[derive(Debug)] 82 | enum SerializedResponse { 83 | Text(String), 84 | Bytes(Vec), 85 | } 86 | 87 | fn parse_response( 88 | prediction: &Value, 89 | schema: &Value, 90 | output: &mut Option, 91 | ) -> SerializedResponse { 92 | if schema.get("type") == Some(&Value::String("array".to_string())) { 93 | todo!("array response not yet supported"); 94 | } 95 | 96 | if schema.get("type") == Some(&Value::String("string".to_string())) 97 | && schema.get("format") == Some(&Value::String("uri".to_string())) 98 | { 99 | let url = prediction.as_str().unwrap(); 100 | 101 | if !url.starts_with("data:") { 102 | return SerializedResponse::Text(url.to_string()); 103 | } 104 | 105 | let dataurl = DataUrl::parse(url).expect("Failed to parse data URI"); 106 | 107 | if output.is_none() { 108 | *output = Some(PathBuf::from(format!( 109 | "output{}", 110 | mime_guess::get_mime_extensions( 111 | &Mime::from_str(dataurl.get_media_type()) 112 | .unwrap_or(mime_guess::mime::APPLICATION_OCTET_STREAM), 113 | ) 114 | .and_then(<[&str]>::last) 115 | .map(|e| format!(".{e}")) 116 | .unwrap_or_default() 117 | ))); 118 | } 119 | 120 | return SerializedResponse::Bytes(dataurl.get_data().to_vec()); 121 | } 122 | 123 | if schema.get("type") == Some(&Value::String("string".to_string())) { 124 | return SerializedResponse::Text( 125 | prediction 126 | .as_str() 127 | .expect("Expected prediction to be a string") 128 | .to_string(), 129 | ); 130 | } 131 | 132 | SerializedResponse::Text(serde_json::to_string(&prediction).unwrap()) 133 | } 134 | 135 | fn parse_inputs(inputs: Vec, schema: &SchemaObject) -> HashMap { 136 | let mut key_vals = HashMap::new(); 137 | 138 | for input in inputs { 139 | let (name, mut value) = if input.contains('=') { 140 | let split: [String; 2] = input 141 | .splitn(2, '=') 142 | .map(str::to_string) 143 | .collect::>() 144 | .try_into() 145 | .expect( 146 | "Failed to parse input. Please specify inputs in the format '-i name=value'", 147 | ); 148 | 149 | (split[0].clone(), split[1].clone()) 150 | } else { 151 | (get_first_input(schema).expect("Could not determine the default input based on the order of the inputs. Please specify inputs in the format '-i name=value'"), input) 152 | }; 153 | 154 | if value.starts_with('"') && value.ends_with('"') { 155 | value = value[1..value.len() - 1].to_string(); 156 | } 157 | 158 | key_vals.insert( 159 | name.clone(), 160 | value.strip_prefix('@').map_or_else( 161 | || value.clone(), 162 | |path_str| { 163 | let bytes = fs::read(PathBuf::from(path_str)) 164 | .expect("Couldn't find {path_str} file (for {name})"); 165 | 166 | let mime = tree_magic_mini::from_u8(bytes.as_slice()); 167 | 168 | format!("data:{mime};base64,{}", Base64.encode(bytes)) 169 | }, 170 | ), 171 | ); 172 | } 173 | 174 | key_vals 175 | } 176 | 177 | pub fn get_first_input(schema: &SchemaObject) -> Option { 178 | let input_properties = schema 179 | .extensions 180 | .get("components")? 181 | .get("schemas")? 182 | .get("Input")? 183 | .get("properties")?; 184 | 185 | for (k, v) in input_properties.as_object()? { 186 | let Some(order) = v.get("x-order").and_then(|o| match o { 187 | Value::Number(n) => n.as_i64(), 188 | _ => None, 189 | }) else { 190 | continue; 191 | }; 192 | 193 | if order == 0 { 194 | return Some(k.clone()); 195 | } 196 | } 197 | 198 | None 199 | } 200 | -------------------------------------------------------------------------------- /cli/src/commands/push.rs: -------------------------------------------------------------------------------- 1 | use crate::Context; 2 | 3 | pub fn handle(ctx: Context, image: &Option) { 4 | let builder = ctx.into_builder(); 5 | 6 | if builder.config.image.as_ref().or(image.as_ref()).is_none() { 7 | eprintln!("To push images, you must either set the 'image' option in the packages.metadata.cog of your Cargo.toml or pass an image name as an argument. For example, 'cargo cog push hotdog-detector'"); 8 | std::process::exit(1); 9 | } 10 | 11 | let image_name = builder.build(image.clone()); 12 | builder.push(image); 13 | 14 | println!("Image '{image_name}' pushed"); 15 | 16 | if image_name.starts_with("r8.im/") { 17 | println!( 18 | "Run your model on Replicate:\n https://{}", 19 | image_name.replacen("r8.im", "replicate.com", 1) 20 | ); 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /cli/src/config.rs: -------------------------------------------------------------------------------- 1 | use cargo_metadata::Package; 2 | use serde::{Deserialize, Serialize}; 3 | use serde_json::json; 4 | 5 | #[derive(Debug, Deserialize, Serialize, Default)] 6 | pub struct Config { 7 | #[serde(default)] 8 | pub gpu: bool, 9 | pub image: Option, 10 | } 11 | 12 | impl Config { 13 | pub fn from_package(package: &Package) -> Self { 14 | let mut config = package 15 | .metadata 16 | .get("cog") 17 | .and_then(|config| Self::deserialize(config).ok()) 18 | .unwrap_or_default(); 19 | 20 | if config.image.is_none() { 21 | config.image = Some(Self::generate_image_name(&package.name)); 22 | } 23 | 24 | config 25 | } 26 | 27 | pub fn image_name(&self, image: Option) -> String { 28 | image.or_else(|| self.image.clone()).unwrap() 29 | } 30 | 31 | fn generate_image_name(name: &str) -> String { 32 | let mut image_name = name 33 | .to_lowercase() 34 | .replace(|c: char| !c.is_alphanumeric(), "-"); 35 | 36 | if !image_name.starts_with("cog-") { 37 | image_name = format!("cog-{image_name}"); 38 | } 39 | 40 | let mut image_name = image_name 41 | .chars() 42 | .take(30 - "cog-".len()) 43 | .collect::(); 44 | 45 | while let Some(last_char) = image_name.chars().last() { 46 | if last_char.is_alphanumeric() { 47 | break; 48 | } 49 | 50 | image_name.pop(); 51 | } 52 | 53 | image_name 54 | } 55 | 56 | #[allow(clippy::unused_self)] 57 | pub fn as_cog_config(&self) -> String { 58 | serde_json::to_string(&json!({ 59 | "predict": "main.rs:CogModel", 60 | "build": { 61 | "gpu": self.gpu, 62 | "python_version" : "N/A" 63 | }, 64 | })) 65 | .unwrap() 66 | } 67 | } 68 | 69 | #[cfg(test)] 70 | mod tests { 71 | use super::*; 72 | 73 | #[test] 74 | fn generate_image_name() { 75 | assert_eq!( 76 | "cog-hello-world", 77 | Config::generate_image_name("hello-world"), 78 | ); 79 | 80 | assert_eq!( 81 | "cog-hello-world", 82 | Config::generate_image_name("cog-hello-world"), 83 | ); 84 | 85 | assert_eq!( 86 | "cog-a-very-very-long-packa", 87 | Config::generate_image_name("a-very-very-long-package-name"), 88 | ); 89 | 90 | assert_eq!( 91 | "cog-with-a-very-very-long", 92 | Config::generate_image_name("cog-with-a-very-very-long-package-name"), 93 | ); 94 | 95 | assert_eq!( 96 | "cog-with-a-very-very-long", 97 | Config::generate_image_name("cog-with-a-very-very-long-package-name"), 98 | ); 99 | 100 | assert_eq!( 101 | "cog-with-invalid-name", 102 | Config::generate_image_name("cog-with-invalid-name-!@#$%^&*()"), 103 | ); 104 | } 105 | } 106 | -------------------------------------------------------------------------------- /cli/src/context.rs: -------------------------------------------------------------------------------- 1 | use std::path::PathBuf; 2 | 3 | use anyhow::Result; 4 | 5 | use crate::docker::Docker; 6 | 7 | #[derive(Debug, Clone)] 8 | pub struct Context { 9 | pub cwd: PathBuf, 10 | } 11 | 12 | impl Context { 13 | /// Create a new context 14 | /// 15 | /// # Errors 16 | /// 17 | /// This function will return an error if the Docker daemon is not running or if the current working directory cannot be determined. 18 | pub fn new() -> Result { 19 | Docker::check_connection()?; 20 | 21 | Ok(Self { 22 | cwd: std::env::current_dir()?, 23 | }) 24 | } 25 | 26 | #[must_use] 27 | pub fn into_builder(self) -> crate::docker::Builder { 28 | crate::docker::Builder::new(self.cwd) 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /cli/src/docker/auth.rs: -------------------------------------------------------------------------------- 1 | use base64::{engine::general_purpose::STANDARD as Base64, Engine}; 2 | use serde_json::{json, Value}; 3 | use std::{ 4 | fs, 5 | io::Write, 6 | process::{Command, Stdio}, 7 | }; 8 | 9 | pub fn store_credentials(registry: &str, username: &str, token: &str) -> Result<(), String> { 10 | let docker_config_path = dirs::home_dir() 11 | .unwrap() 12 | .join(".docker") 13 | .join("config.json"); 14 | 15 | if !docker_config_path.exists() { 16 | return Err(format!( 17 | "Couldn't find Docker config file at {}", 18 | docker_config_path.display() 19 | )); 20 | } 21 | 22 | let mut docker_config: Value = 23 | serde_json::from_str(&fs::read_to_string(&docker_config_path).unwrap()).unwrap(); 24 | let credential_store = docker_config 25 | .get_mut("credsStore") 26 | .and_then(|v| Some(v.as_str()?.to_string())); 27 | 28 | if let Some(credential_store) = credential_store { 29 | save_in_store(&credential_store, registry, username, token)?; 30 | } else { 31 | save_in_config(&mut docker_config, registry, username, token); 32 | 33 | fs::write( 34 | docker_config_path, 35 | serde_json::to_string_pretty(&docker_config).unwrap(), 36 | ) 37 | .expect("Failed to save Docker config."); 38 | } 39 | 40 | Ok(()) 41 | } 42 | 43 | fn save_in_store(store: &str, registry: &str, username: &str, token: &str) -> Result<(), String> { 44 | let binary = format!("docker-credential-{store}"); 45 | 46 | let mut cmd = Command::new(&binary).stdin(Stdio::piped()).spawn().unwrap(); 47 | 48 | let stdin = cmd.stdin.as_mut().unwrap(); 49 | stdin 50 | .write_all( 51 | json!({ "ServerURL": registry, "Username": username, "Secret": token }) 52 | .to_string() 53 | .as_bytes(), 54 | ) 55 | .unwrap(); 56 | 57 | let output = cmd.wait_with_output().unwrap(); 58 | 59 | if !output.status.success() { 60 | return Err(format!( 61 | "Failed to store credentials using {}: {}", 62 | binary, 63 | String::from_utf8_lossy(&output.stderr) 64 | )); 65 | } 66 | 67 | Ok(()) 68 | } 69 | 70 | fn save_in_config(docker_config: &mut Value, registry: &str, username: &str, token: &str) { 71 | let auths = docker_config 72 | .get_mut("auths") 73 | .unwrap() 74 | .as_object_mut() 75 | .unwrap(); 76 | 77 | auths.insert( 78 | registry.to_string(), 79 | serde_json::json!({ "auth": Base64.encode(format!("{username}:{token}")) }), 80 | ); 81 | } 82 | -------------------------------------------------------------------------------- /cli/src/docker/builder.rs: -------------------------------------------------------------------------------- 1 | use cargo_metadata::{MetadataCommand, Package}; 2 | use map_macro::hash_map; 3 | use std::{ 4 | collections::HashMap, 5 | fs::{self, File}, 6 | io::Write, 7 | path::PathBuf, 8 | process::{Command, Stdio}, 9 | }; 10 | 11 | use super::dockerfile::{Dockerfile, DockerfileExt}; 12 | use crate::{ 13 | config::Config, 14 | docker::{Docker, Error as DockerError, RunOptions}, 15 | }; 16 | 17 | pub struct Builder { 18 | cwd: PathBuf, 19 | package: Package, 20 | pub config: Config, 21 | _cog_version: String, 22 | deps: Vec, 23 | } 24 | 25 | impl Builder { 26 | pub fn new(cwd: PathBuf) -> Self { 27 | let cargo_metadata = MetadataCommand::new() 28 | .manifest_path(cwd.join("Cargo.toml")) 29 | .exec() 30 | .expect( 31 | "Failed to read Cargo.toml. Make sure you are in the root of your Cog project.", 32 | ); 33 | 34 | let package = cargo_metadata 35 | .root_package() 36 | .expect("Couldn't find the package section in Cargo.toml."); 37 | 38 | assert!( 39 | !package.authors.is_empty(), 40 | "You must specify at least one author in Cargo.toml" 41 | ); 42 | 43 | let cog_version = package 44 | .dependencies 45 | .iter() 46 | .find(|dep| dep.name == "cog-rust") 47 | .expect("Couldn't find cog-rust in your Cargo.toml") 48 | .req 49 | .to_string(); 50 | 51 | assert!(cog_version != "*", "Couldn't resolve cog version. Make sure you're loading the package through the registry, not from git or a local path."); 52 | 53 | Self { 54 | cwd, 55 | package: package.clone(), 56 | _cog_version: cog_version, 57 | deps: cargo_metadata.packages.clone(), 58 | config: Config::from_package(package), 59 | } 60 | } 61 | 62 | pub fn generate_dockerfile(&self) -> String { 63 | let torchlib = || { 64 | self.deps.iter().find(|dep| dep.name == "torch-sys")?; 65 | 66 | let pytorch_url = if self.config.gpu { 67 | "https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu118.zip" 68 | } else { 69 | "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcpu.zip" 70 | }; 71 | 72 | Some( 73 | Dockerfile::new() 74 | .run_multiple(&[ 75 | Command::new("curl").args(["-sSL", pytorch_url, "-o libtorch.zip"]), 76 | Command::new("unzip").arg("libtorch.zip"), 77 | Command::new("rm").arg("libtorch.zip"), 78 | Command::new("cp").arg("libtorch/lib/*").arg("/src/lib"), 79 | ]) 80 | .env("LIBTORCH", "/src/libtorch"), 81 | ) 82 | }; 83 | 84 | let weights_dir = || { 85 | if self.cwd.join("weights").exists() { 86 | Some(Dockerfile::new().copy("weights/", "/src/weights/")) 87 | } else { 88 | None 89 | } 90 | }; 91 | 92 | let replicate_hack = || { 93 | Some( 94 | Dockerfile::new() 95 | .run_multiple(&[ 96 | Command::new("echo") 97 | .arg(r##""#!/bin/bash\nexit 0""##) 98 | .arg(">") 99 | .arg("/usr/local/bin/pip"), 100 | Command::new("chmod").arg("+x").arg("/usr/local/bin/pip"), 101 | Command::new("echo") 102 | .arg(format!( 103 | // Replicate runs `python -m cog.server.http --other-args*` and we only care about the other args 104 | "\"#!/bin/bash\\nshift 2; /usr/bin/{} \"\\$@\"\"", 105 | self.package.name 106 | )) 107 | .arg(">") 108 | .arg("/usr/local/bin/python"), 109 | Command::new("chmod").arg("+x").arg("/usr/local/bin/python"), 110 | ]) 111 | // Replicate also doesn't provide a way to set the log level, so we have to do it manually 112 | .env("RUST_LOG", r#""cog_rust=trace""#), 113 | ) 114 | }; 115 | 116 | include_str!("../templates/Dockerfile") 117 | .to_string() 118 | .for_bin(&self.package.name) 119 | .handler("before_build", torchlib) 120 | .handler("before_runtime", weights_dir) 121 | .handler("after_runtime", replicate_hack) 122 | .into_image(if self.config.gpu { 123 | "nvidia/cuda:12.2.0-base-ubuntu22.04" 124 | } else { 125 | "debian:bookworm-slim" 126 | }) 127 | .build() 128 | } 129 | 130 | pub fn build(&self, tag: Option) -> String { 131 | let dockerfile = self.generate_dockerfile(); 132 | 133 | File::create(self.cwd.join(".dockerignore")).and_then(|mut file| write!(file, "target")).expect( 134 | "Failed to create .dockerignore file. Make sure you are in the root of your Cog project." 135 | ); 136 | 137 | let image_name = self.config.image_name(tag); 138 | Self::build_image(&dockerfile, &image_name, None, true); 139 | 140 | fs::remove_file(self.cwd.join(".dockerignore")).expect("Failed to clean up .dockerignore"); 141 | 142 | println!("Adding labels to image..."); 143 | let schema = match Docker::run(RunOptions { 144 | image: image_name.clone(), 145 | env: vec!["RUST_LOG=cog_rust=error".to_string()], 146 | flags: vec!["--dump-schema-and-exit".to_string()], 147 | cmd: Some(format!("/usr/bin/{}", self.package.name)), 148 | ..RunOptions::default() 149 | }) { 150 | Ok(output) => output, 151 | Err(DockerError::Parse(_)) => panic!("Failed to parse schema."), 152 | Err(DockerError::Command(err)) => panic!("Failed to extract schema from image: {err}"), 153 | _ => panic!("Failed to extract schema from image."), 154 | }; 155 | 156 | Self::build_image( 157 | &format!("FROM {image_name}"), 158 | &image_name, 159 | Some(hash_map! { 160 | "run.cog.has_init" => "true", 161 | // Seems like Replicate will only allow `cog` versions, so we hardcode it here 162 | "run.cog.version" => "0.7.2", 163 | "org.cogmodel.cog_version" => "0.7.2", 164 | "run.cog.openapi_schema" => schema.trim(), 165 | "org.cogmodel.openapi_schema" => schema.trim(), 166 | "run.cog.config" => &self.config.as_cog_config(), 167 | "rs.cog.authors" => &self.package.authors.join(", "), 168 | "org.cogmodel.config" => &self.config.as_cog_config(), 169 | "org.cogmodel.deprecated" => "The org.cogmodel labels are deprecated. Use run.cog.", 170 | }), 171 | false, 172 | ); 173 | 174 | image_name 175 | } 176 | 177 | pub fn push(&self, image: &Option) { 178 | Docker::push( 179 | image 180 | .as_ref() 181 | .or(self.config.image.as_ref()) 182 | .expect("Image name not specified"), 183 | ) 184 | .expect("Failed to push image."); 185 | } 186 | 187 | fn build_image( 188 | dockerfile: &str, 189 | image_name: &str, 190 | labels: Option>, 191 | show_logs: bool, 192 | ) { 193 | let mut process = Command::new("docker") 194 | .arg("build") 195 | .args([ 196 | "--platform", 197 | "linux/amd64", 198 | "--tag", 199 | image_name, 200 | "--build-arg", 201 | "BUILDKIT_INLINE_CACHE=1", 202 | "--file", 203 | "-", 204 | ]) 205 | .args( 206 | labels 207 | .unwrap_or_default() 208 | .iter() 209 | .flat_map(|(key, value)| ["--label".to_string(), format!("{key}={value}")]), 210 | ) 211 | .arg(".") 212 | .env("DOCKER_BUILDKIT", "1") 213 | .env("DOCKER_DEFAULT_PLATFORM", "linux/amd64") 214 | .stdin(Stdio::piped()) 215 | .stderr(if show_logs { 216 | Stdio::inherit() 217 | } else { 218 | Stdio::null() 219 | }) 220 | .stdout(if show_logs { 221 | Stdio::inherit() 222 | } else { 223 | Stdio::null() 224 | }) 225 | .spawn() 226 | .expect("Failed to spawn docker build process."); 227 | 228 | process 229 | .stdin 230 | .as_mut() 231 | .expect("Failed to open stdin.") 232 | .write_all(dockerfile.as_bytes()) 233 | .expect("Failed to write to stdin."); 234 | 235 | let status = process 236 | .wait() 237 | .expect("Failed to wait for docker build process."); 238 | 239 | assert!(status.success(), "Failed to build docker image."); 240 | } 241 | } 242 | 243 | impl Drop for Builder { 244 | fn drop(&mut self) { 245 | let _ = fs::remove_file(self.cwd.join(".dockerignore")); 246 | } 247 | } 248 | -------------------------------------------------------------------------------- /cli/src/docker/dockerfile.rs: -------------------------------------------------------------------------------- 1 | use std::process::Command; 2 | 3 | #[derive(Debug)] 4 | pub struct Dockerfile(String); 5 | 6 | impl Dockerfile { 7 | pub const fn new() -> Self { 8 | Self(String::new()) 9 | } 10 | 11 | pub fn run_multiple(mut self, command: &[&Command]) -> Self { 12 | self.0.push_str(&format!( 13 | "RUN {}\n", 14 | command 15 | .iter() 16 | .map(|cmd| format!( 17 | "{} {}", 18 | cmd.get_program().to_string_lossy(), 19 | cmd.get_args() 20 | .map(|arg| arg.to_string_lossy().to_string()) 21 | .collect::>() 22 | .join(" ") 23 | )) 24 | .collect::>() 25 | .join(" && ") 26 | )); 27 | 28 | self 29 | } 30 | 31 | pub fn env(mut self, key: &str, value: &str) -> Self { 32 | self.0.push_str(&format!("ENV {key}={value}\n")); 33 | 34 | self 35 | } 36 | 37 | pub fn copy(mut self, src: &str, dest: &str) -> Self { 38 | self.0.push_str(&format!("COPY {src} {dest}\n")); 39 | 40 | self 41 | } 42 | } 43 | 44 | #[allow(clippy::module_name_repetitions)] 45 | pub trait DockerfileExt { 46 | fn build(self) -> Self; 47 | fn for_bin(self, bin: &str) -> Self; 48 | fn into_image(self, image: &str) -> Self; 49 | fn handler(self, slot: &str, value: impl FnOnce() -> Option) -> Self; 50 | } 51 | 52 | impl DockerfileExt for String { 53 | fn for_bin(self, bin: &str) -> Self { 54 | self.replace("{:bin_name}", bin) 55 | } 56 | 57 | fn into_image(self, image: &str) -> Self { 58 | self.replace("{:base_image}", image) 59 | } 60 | 61 | fn handler(self, slot: &str, value: impl FnOnce() -> Option) -> Self { 62 | let value = value(); 63 | 64 | match value { 65 | Some(value) => self.replace( 66 | &format!("#SLOT {slot}"), 67 | &format!("{}#SLOT {slot}\n", value.0), 68 | ), 69 | None => self, 70 | } 71 | } 72 | 73 | fn build(self) -> Self { 74 | self.lines() 75 | .filter(|line| !line.starts_with("#SLOT")) 76 | .collect::>() 77 | .join("\n") 78 | .replace("\n\n", "\n") 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /cli/src/docker/mod.rs: -------------------------------------------------------------------------------- 1 | mod auth; 2 | mod builder; 3 | mod dockerfile; 4 | pub mod predictor; 5 | 6 | use std::{ 7 | collections::HashMap, 8 | io, 9 | path::PathBuf, 10 | process::{self, Command, Stdio}, 11 | }; 12 | 13 | pub use auth::store_credentials; 14 | pub use builder::Builder; 15 | pub use predictor::Predictor; 16 | 17 | /// Errors that can occur when interacting with the docker CLI. 18 | #[derive(Debug, thiserror::Error)] 19 | pub enum Error { 20 | #[error("Could not connect to Docker. Is the docker daemon running?")] 21 | NotRunning, 22 | 23 | #[error("The provided image could not be found.")] 24 | NotFound, 25 | 26 | #[error("Provided flags without a command.")] 27 | CmdMissing, 28 | 29 | #[error("{0}")] 30 | Command(String), 31 | 32 | #[error("Failed to parse output from command: {0}")] 33 | Parse(String), 34 | 35 | #[error("Failed to run command: {0}")] 36 | Spawn(#[from] std::io::Error), 37 | 38 | #[error("Failed to parse output from command: {0}")] 39 | ToString(#[from] std::string::FromUtf8Error), 40 | 41 | #[error("Failed to parse output from command: {0}")] 42 | Deserialize(#[from] serde_json::Error), 43 | } 44 | 45 | #[derive(Debug, Default)] 46 | pub struct RunOptions { 47 | detach: bool, 48 | image: String, 49 | env: Vec, 50 | interactive: bool, 51 | flags: Vec, 52 | cmd: Option, 53 | ports: HashMap, 54 | volumes: HashMap, 55 | } 56 | 57 | /// A wrapper around the docker CLI. 58 | pub struct Docker {} 59 | 60 | impl Docker { 61 | /// Check if the docker daemon is running. 62 | /// 63 | /// # Errors 64 | /// 65 | /// Returns an error if the docker daemon is not running. 66 | pub fn check_connection() -> Result<(), Error> { 67 | let status = Command::new("docker") 68 | .arg("info") 69 | .stdout(Stdio::null()) 70 | .status()?; 71 | 72 | if !status.success() { 73 | return Err(Error::NotRunning); 74 | } 75 | 76 | Ok(()) 77 | } 78 | 79 | /// Inspect the given image. 80 | /// Returns the image metadata as a JSON struct. 81 | /// 82 | /// # Errors 83 | /// 84 | /// Returns an error if the image could not be found. 85 | pub fn inspect_image(image: &str) -> Result { 86 | let output = Command::new("docker") 87 | .arg("image") 88 | .arg("inspect") 89 | .arg(image) 90 | .output()?; 91 | 92 | if !output.status.success() 93 | && String::from_utf8(output.stderr.clone())?.contains("No such image") 94 | { 95 | return Err(Error::NotFound); 96 | } 97 | 98 | if !output.status.success() { 99 | return Err(Error::Command(format!( 100 | "Failed to inspect image: {}", 101 | String::from_utf8(output.stderr)?.trim() 102 | ))); 103 | } 104 | 105 | Ok(serde_json::from_slice(&output.stdout)?) 106 | } 107 | 108 | /// Inspect the given container. 109 | /// Returns the container metadata as a JSON struct. 110 | /// 111 | /// # Errors 112 | /// 113 | /// Returns an error if the container could not be found. 114 | pub fn inspect_container(container_id: &str) -> Result { 115 | let output = Command::new("docker") 116 | .arg("container") 117 | .arg("inspect") 118 | .arg(container_id) 119 | .output()?; 120 | 121 | if !output.status.success() 122 | && String::from_utf8(output.stderr.clone())?.contains("No such container") 123 | { 124 | return Err(Error::NotFound); 125 | } 126 | 127 | if !output.status.success() { 128 | return Err(Error::Command(format!( 129 | "Failed to inspect container: {}", 130 | String::from_utf8(output.stderr)?.trim() 131 | ))); 132 | } 133 | 134 | Ok(serde_json::from_slice(&output.stdout)?) 135 | } 136 | 137 | /// Pull the given image from the Docker registry. 138 | /// Returns the image digest. 139 | /// 140 | /// # Errors 141 | /// 142 | /// Returns an error if the command fails or if the output cannot be parsed. 143 | pub fn pull(image: &str) -> Result { 144 | let output = Command::new("docker") 145 | .arg("pull") 146 | .arg(image) 147 | .stdout(Stdio::inherit()) 148 | .stderr(Stdio::inherit()) 149 | .output()?; 150 | 151 | if !output.status.success() { 152 | return Err(Error::Command(format!( 153 | "Failed to pull image: {}", 154 | String::from_utf8(output.stderr)?.trim() 155 | ))); 156 | } 157 | 158 | Ok(String::from_utf8(output.stdout)? 159 | .split("sha256:") 160 | .last() 161 | .ok_or_else(|| Error::Parse("docker pull".to_string()))? 162 | .split(' ') 163 | .next() 164 | .ok_or_else(|| Error::Parse("docker pull".to_string()))? 165 | .to_string()) 166 | } 167 | 168 | /// Push the given image to the Docker registry. 169 | /// Returns the image digest. 170 | /// 171 | /// # Errors 172 | /// 173 | /// Returns an error if the command fails or if the output cannot be parsed. 174 | pub fn push(image: &str) -> Result { 175 | let output = Command::new("docker") 176 | .arg("push") 177 | .arg(image) 178 | .stdout(Stdio::inherit()) 179 | .stderr(Stdio::inherit()) 180 | .output()?; 181 | 182 | if !output.status.success() { 183 | return Err(Error::Command(format!( 184 | "Failed to push image: {}", 185 | String::from_utf8(output.stderr)?.trim() 186 | ))); 187 | } 188 | 189 | Ok(String::from_utf8(output.stdout)? 190 | .split("sha256:") 191 | .last() 192 | .ok_or_else(|| Error::Parse("docker push".to_string()))? 193 | .split(' ') 194 | .next() 195 | .ok_or_else(|| Error::Parse("docker push".to_string()))? 196 | .to_string()) 197 | } 198 | 199 | /// Run a container with the given image, volumes, and ports. 200 | /// Returns the container ID. 201 | /// 202 | /// # Errors 203 | /// 204 | /// Returns an error if the command fails or if the output cannot be parsed. 205 | pub fn run(opts: RunOptions) -> Result { 206 | let mut cmd = Command::new("docker"); 207 | 208 | cmd.arg("run").arg("--rm").stderr(Stdio::piped()); 209 | 210 | if opts.detach { 211 | cmd.arg("--detach"); 212 | } 213 | 214 | if opts.interactive { 215 | cmd.arg("--interactive"); 216 | } 217 | 218 | for var in opts.env { 219 | cmd.args(["--env", &var]); 220 | } 221 | 222 | for (source, destination) in opts.ports { 223 | cmd.args(["--publish", &format!("{source}:{destination}")]); 224 | } 225 | 226 | for (source, destination) in opts.volumes { 227 | cmd.args([ 228 | "--mount", 229 | &format!( 230 | "type=bind,source={},destination={destination}", 231 | source.display() 232 | ), 233 | ]); 234 | } 235 | 236 | cmd.arg(opts.image); 237 | 238 | if let Some(bin) = opts.cmd { 239 | cmd.arg(bin); 240 | 241 | for flag in opts.flags { 242 | cmd.arg(flag); 243 | } 244 | } else if !opts.flags.is_empty() { 245 | return Err(Error::CmdMissing); 246 | } 247 | 248 | let output = cmd.stderr(Stdio::piped()).output()?; 249 | 250 | if !output.status.success() { 251 | return Err(Error::Command(format!( 252 | "Failed to start container: {}", 253 | String::from_utf8(output.stderr)?.trim() 254 | ))); 255 | } 256 | 257 | Ok(String::from_utf8(output.stdout) 258 | .map_err(|_| Error::Parse("docker run".to_string()))? 259 | .trim() 260 | .to_string()) 261 | } 262 | 263 | pub fn stop(container_id: &str) -> Result<(), Error> { 264 | let output = Command::new("docker") 265 | .arg("container") 266 | .arg("stop") 267 | .args(["--time", "3"]) 268 | .arg(container_id) 269 | .stderr(Stdio::inherit()) 270 | .output()?; 271 | 272 | if !output.status.success() { 273 | return Err(Error::Command(format!( 274 | "Failed to stop container: {}", 275 | String::from_utf8(output.stderr)?.trim() 276 | ))); 277 | } 278 | 279 | Ok(()) 280 | } 281 | 282 | pub fn find_port(container_id: &str, container_port: u16) -> Result { 283 | let output = Command::new("docker") 284 | .arg("port") 285 | .arg(container_id) 286 | .arg(container_port.to_string()) 287 | .output()?; 288 | 289 | if !output.status.success() { 290 | return Err(Error::Command(format!( 291 | "Failed to get port: {}", 292 | String::from_utf8(output.stderr)?.trim() 293 | ))); 294 | } 295 | 296 | let stdout = String::from_utf8(output.stdout)?; 297 | 298 | stdout 299 | .trim() 300 | .split(':') 301 | .last() 302 | .ok_or_else(|| Error::Parse("docker port".to_string()))? 303 | .parse() 304 | .map_err(|_| Error::Parse("docker port".to_string())) 305 | } 306 | 307 | pub fn tail_logs(container_id: &str) -> io::Result { 308 | Command::new("docker") 309 | .arg("container") 310 | .arg("logs") 311 | .arg("-f") 312 | .arg(container_id) 313 | .stdout(Stdio::inherit()) 314 | .stderr(Stdio::inherit()) 315 | .spawn() 316 | } 317 | } 318 | -------------------------------------------------------------------------------- /cli/src/docker/predictor.rs: -------------------------------------------------------------------------------- 1 | use anyhow::{bail, Context}; 2 | use cog_core::http::{HTTPValidationError, Response}; 3 | use indoc::formatdoc; 4 | use map_macro::hash_map; 5 | use reqwest::StatusCode; 6 | use schemars::schema::SchemaObject; 7 | use serde_json::{json, Value}; 8 | use std::{ 9 | collections::HashMap, 10 | time::{Duration, Instant}, 11 | }; 12 | use tokio::time::sleep; 13 | 14 | use super::RunOptions; 15 | use crate::docker::Docker; 16 | 17 | #[derive(Debug)] 18 | pub struct Predictor { 19 | image: String, 20 | port: Option, 21 | container_id: Option, 22 | } 23 | 24 | impl Predictor { 25 | pub const fn new(image: String) -> Self { 26 | Self { 27 | image, 28 | port: None, 29 | container_id: None, 30 | } 31 | } 32 | 33 | pub fn get_schema(&self) -> anyhow::Result { 34 | let image = Docker::inspect_image(&self.image)?; 35 | 36 | serde_json::from_str::( 37 | image 38 | .as_array() 39 | .and_then(|v| v.first()) 40 | .and_then(|v| { 41 | v.get("Config") 42 | .and_then(Value::as_object) 43 | .and_then(|v| v.get("Labels")) 44 | .and_then(Value::as_object) 45 | .and_then(|v| v.get("org.cogmodel.openapi_schema")) 46 | .and_then(Value::as_str) 47 | }) 48 | .context("Failed to get schema label. Is this a Cog model?")?, 49 | ) 50 | .context("Failed to parse schema") 51 | } 52 | 53 | pub async fn start(&mut self) { 54 | let container_id = Docker::run(RunOptions { 55 | detach: true, 56 | image: self.image.clone(), 57 | ports: hash_map! { 0 => 5000 }, 58 | env: vec!["RUST_LOG=error".to_string()], 59 | ..RunOptions::default() 60 | }) 61 | .unwrap(); 62 | 63 | self.container_id = Some(container_id.clone()); 64 | self.port = Some(Docker::find_port(&container_id, 5000).unwrap()); 65 | 66 | Docker::tail_logs(&container_id).unwrap(); 67 | self.wait_for_server().await.unwrap(); 68 | } 69 | 70 | pub async fn predict(&self, inputs: HashMap) -> anyhow::Result { 71 | let port = self 72 | .port 73 | .as_ref() 74 | .context("Trying to predict with non-running container.")?; 75 | 76 | let client = reqwest::Client::new(); 77 | 78 | let res = client 79 | .post(format!("http://localhost:{port}/predictions")) 80 | .json(&json!({ "input": inputs })) 81 | .send() 82 | .await 83 | .context("Failed to send request")?; 84 | 85 | if matches!(res.status(), StatusCode::UNPROCESSABLE_ENTITY) { 86 | let text = res.text().await?; 87 | let errors = serde_json::from_str::(&text).context( 88 | format!("/predictions call returned status 422, and the response body failed to decode: {text}") 89 | )?; 90 | 91 | bail!(formatdoc! {" 92 | The inputs you passed to cog predict could not be validated: 93 | 94 | {} 95 | 96 | You can provide an input with -i. For example: 97 | 98 | cog predict -i blur=3.5 99 | 100 | If your input is a local file, you need to prefix the path with @ to tell Cog to read the file contents. For example: 101 | 102 | cog predict -i path=@image.jpg 103 | ", errors.detail.iter().map(|e| e.msg.clone()).collect::>().join("\n")}); 104 | } 105 | 106 | if !matches!(res.status(), StatusCode::OK) { 107 | bail!("/predictions call returned status {}", res.status()); 108 | } 109 | 110 | let text = res.text().await?; 111 | serde_json::from_str::(&text) 112 | .context("Failed to decode prediction response: {text}") 113 | } 114 | 115 | pub fn stop(&self) -> anyhow::Result<()> { 116 | let container_id = self 117 | .container_id 118 | .as_ref() 119 | .context("Trying to stop non-running container.")?; 120 | 121 | Ok(Docker::stop(container_id)?) 122 | } 123 | 124 | pub async fn wait_for_server(&self) -> anyhow::Result<()> { 125 | let start = Instant::now(); 126 | 127 | let container_id = self 128 | .container_id 129 | .as_ref() 130 | .context("Waiting for non-running container.")?; 131 | 132 | let client = reqwest::Client::new(); 133 | 134 | loop { 135 | if start.elapsed().as_secs() > 300 { 136 | return Err(anyhow::anyhow!("Timed out")); 137 | } 138 | 139 | sleep(Duration::from_millis(100)).await; 140 | 141 | let container = Docker::inspect_container(container_id) 142 | .context("Failed to get container status")?; 143 | 144 | let state = container 145 | .as_array() 146 | .and_then(|v| v.first()) 147 | .and_then(Value::as_object) 148 | .and_then(|v| v.get("State")) 149 | .and_then(Value::as_object) 150 | .and_then(|v| v.get("Status")) 151 | .and_then(Value::as_str) 152 | .expect("Container exited unexpectedly"); 153 | 154 | if state == "exited" || state == "dead" { 155 | bail!("Container exited unexpectedly"); 156 | } 157 | 158 | let res = client 159 | .get(format!( 160 | "http://localhost:{}/health-check", 161 | self.port.unwrap() 162 | )) 163 | .send() 164 | .await 165 | .and_then(reqwest::Response::error_for_status); 166 | 167 | let Ok(res) = res else { 168 | continue; 169 | }; 170 | 171 | let status = res 172 | .json::() 173 | .await 174 | .ok() 175 | .and_then(|v| v.get("status").cloned()) 176 | .and_then(|s| s.as_str().map(str::to_string)) 177 | .context("Container healthcheck returned invalid response")?; 178 | 179 | match status.as_str() { 180 | "STARTING" => continue, 181 | "READY" => return Ok(()), 182 | "SETUP_FAILED" => bail!("Model setup failed"), 183 | _ => bail!("Container healthcheck returned unexpected status: {status}"), 184 | } 185 | } 186 | } 187 | } 188 | 189 | impl Drop for Predictor { 190 | fn drop(&mut self) { 191 | let _ = self.stop(); 192 | } 193 | } 194 | -------------------------------------------------------------------------------- /cli/src/helpers.rs: -------------------------------------------------------------------------------- 1 | use std::io::{self, Read}; 2 | 3 | #[must_use] 4 | pub fn wait_for_input() -> String { 5 | let mut input = String::new(); 6 | io::stdin() 7 | .read_line(&mut input) 8 | .expect("Failed to read line"); 9 | 10 | input.trim().to_string() 11 | } 12 | 13 | #[must_use] 14 | pub fn load_from_stdin() -> String { 15 | let mut input = String::new(); 16 | io::stdin() 17 | .read_to_string(&mut input) 18 | .expect("Failed to load from stdin"); 19 | 20 | input.trim().to_string() 21 | } 22 | -------------------------------------------------------------------------------- /cli/src/main.rs: -------------------------------------------------------------------------------- 1 | #![warn(clippy::all, clippy::pedantic, clippy::nursery)] 2 | 3 | use clap::Parser; 4 | use context::Context; 5 | use CargoSubcommand::Cog; 6 | 7 | mod commands; 8 | mod config; 9 | mod context; 10 | mod docker; 11 | mod helpers; 12 | 13 | /// Cog's CLI interface 14 | /// 15 | /// This binary should be invoked by Cargo with the new `cog` subcommand. If 16 | /// you're reading this, consider manually adding `cog` as the first argument. 17 | #[derive(Debug, Parser)] 18 | struct Cargo { 19 | #[clap(subcommand)] 20 | command: CargoSubcommand, 21 | } 22 | 23 | #[derive(Debug, Parser)] 24 | pub enum CargoSubcommand { 25 | /// A cargo subcommand to build, run and publish machine learning containers 26 | Cog(Cli), 27 | } 28 | 29 | #[derive(Parser, Debug)] 30 | #[command(about, author, display_name = "cargo-cog")] 31 | #[command(override_usage = "cargo cog [OPTIONS] [COMMAND]")] 32 | pub struct Cli { 33 | #[command(subcommand)] 34 | pub command: commands::Command, 35 | } 36 | 37 | #[tokio::main] 38 | async fn main() { 39 | let cargo = Cargo::parse(); 40 | let Cog(cli) = cargo.command; 41 | 42 | commands::exec(cli.command, Context::new().unwrap()).await; 43 | } 44 | -------------------------------------------------------------------------------- /cli/src/templates/Dockerfile: -------------------------------------------------------------------------------- 1 | # syntax = docker/dockerfile:1.2 2 | FROM lukemathwalker/cargo-chef:latest-rust-bookworm AS chef 3 | WORKDIR /src 4 | #SLOT setup 5 | 6 | FROM chef AS planner 7 | COPY . . 8 | RUN cargo chef prepare --bin {:bin_name} --recipe-path recipe.json 9 | 10 | FROM chef AS builder 11 | RUN mkdir /src/lib 12 | RUN apt-get update && apt-get install -y cmake clang 13 | RUN curl -L --proto '=https' --tlsv1.2 -sSf "https://github.com/cargo-bins/cargo-quickinstall/releases/download/cargo-deb-1.44.0/cargo-deb-1.44.0-x86_64-unknown-linux-gnu.tar.gz" | tar -xzvvf - -C /usr/local/cargo/bin 14 | #SLOT before_build 15 | COPY --from=planner /src/recipe.json recipe.json 16 | RUN cargo chef cook --release --bin {:bin_name} 17 | COPY . . 18 | RUN cargo deb --output /src/target/{:bin_name}.deb 19 | #SLOT after_build 20 | 21 | FROM {:base_image} as runtime 22 | WORKDIR /src 23 | RUN apt-get update && apt-get install -y ca-certificates 24 | COPY --from=builder /src/lib/* /lib/x86_64-linux-gnu/ 25 | COPY --from=builder /src/target/{:bin_name}.deb /src/{:bin_name}.deb 26 | #SLOT before_runtime 27 | RUN apt-get install -y /src/{:bin_name}.deb && rm /src/{:bin_name}.deb 28 | EXPOSE 5000 29 | CMD ["/usr/bin/{:bin_name}"] 30 | #SLOT after_runtime 31 | -------------------------------------------------------------------------------- /core/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-core" 3 | version = "0.2.0" 4 | description = "Core types and traits for rust-cog, a Rust toolkit for machine learning." 5 | readme = { workspace = true } 6 | edition = { workspace = true } 7 | authors = { workspace = true } 8 | license = { workspace = true } 9 | keywords = { workspace = true } 10 | categories = { workspace = true } 11 | repository = { workspace = true } 12 | 13 | [dependencies] 14 | anyhow = "1.0.71" 15 | serde = "1.0.164" 16 | thiserror = "1.0.40" 17 | serde_json = "1.0.96" 18 | url = { version = "2.4.0", features = ["serde"] } 19 | tokio = { version = "1.31.0", features = ["rt"] } 20 | chrono = { version = "0.4.26", features = ["serde"] } 21 | schemars = { version = "0.8.12", features = ["url", "chrono"] } 22 | -------------------------------------------------------------------------------- /core/src/http.rs: -------------------------------------------------------------------------------- 1 | use chrono::{DateTime, Utc}; 2 | use schemars::JsonSchema; 3 | use serde::{Deserialize, Serialize}; 4 | use serde_json::Value; 5 | use std::collections::HashMap; 6 | use url::Url; 7 | 8 | #[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema)] 9 | #[serde(rename_all = "lowercase")] 10 | pub enum Status { 11 | #[serde(skip)] 12 | Idle, 13 | 14 | Failed, 15 | Starting, 16 | Canceled, 17 | Succeeded, 18 | Processing, 19 | } 20 | 21 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, JsonSchema)] 22 | pub enum WebhookEvent { 23 | Start, 24 | Output, 25 | Logs, 26 | Completed, 27 | } 28 | 29 | #[derive(Debug, Clone, serde::Deserialize, JsonSchema)] 30 | pub struct Request { 31 | pub webhook: Option, 32 | pub webhook_event_filters: Option>, 33 | 34 | pub input: T, 35 | } 36 | 37 | #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] 38 | pub struct Response { 39 | pub input: Option, 40 | pub output: Option, 41 | 42 | pub id: Option, 43 | pub version: Option, 44 | 45 | pub created_at: Option>, 46 | pub started_at: Option>, 47 | pub completed_at: Option>, 48 | 49 | pub logs: String, 50 | pub status: Status, 51 | pub error: Option, 52 | 53 | pub metrics: Option>, 54 | } 55 | 56 | impl Default for Response { 57 | fn default() -> Self { 58 | Self { 59 | id: None, 60 | error: None, 61 | input: None, 62 | output: None, 63 | metrics: None, 64 | version: None, 65 | created_at: None, 66 | logs: String::new(), 67 | status: Status::Starting, 68 | started_at: Utc::now().into(), 69 | completed_at: Utc::now().into(), 70 | } 71 | } 72 | } 73 | 74 | #[derive(Debug, Deserialize, Serialize)] 75 | pub struct HTTPValidationError { 76 | pub detail: Vec, 77 | } 78 | 79 | #[derive(Debug, Clone, Deserialize, Serialize)] 80 | pub struct ValidationError { 81 | pub msg: String, 82 | pub loc: Vec, 83 | } 84 | -------------------------------------------------------------------------------- /core/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![warn(clippy::all, clippy::pedantic, clippy::nursery)] 2 | 3 | pub mod http; 4 | mod spec; 5 | 6 | pub use spec::{Cog, CogResponse}; 7 | -------------------------------------------------------------------------------- /core/src/spec.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use core::fmt::Debug; 3 | use schemars::JsonSchema; 4 | use serde::{de::DeserializeOwned, Serialize}; 5 | use serde_json::Value; 6 | use std::future::Future; 7 | 8 | use crate::http::Request; 9 | 10 | /// A Cog model 11 | pub trait Cog: Sized + Send { 12 | type Request: DeserializeOwned + JsonSchema + Send; 13 | type Response: CogResponse + Debug + JsonSchema; 14 | 15 | /// Setup the model 16 | /// 17 | /// # Errors 18 | /// 19 | /// Returns an error if setup fails. 20 | fn setup() -> impl Future> + Send; 21 | 22 | /// Run a prediction on the model 23 | /// 24 | /// # Errors 25 | /// 26 | /// Returns an error if the prediction fails. 27 | fn predict(&self, input: Self::Request) -> Result; 28 | } 29 | 30 | /// A response from a Cog model 31 | pub trait CogResponse: Send { 32 | /// Convert the response into a JSON value 33 | fn into_response(self, request: Request) -> impl Future> + Send; 34 | } 35 | 36 | impl CogResponse for T { 37 | async fn into_response(self, _: Request) -> Result { 38 | // We use spawn_blocking here to allow blocking code in serde Serialize impls (used in `Path`, for example). 39 | Ok(tokio::task::spawn_blocking(move || serde_json::to_value(self)).await??) 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /examples/blur/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-blur" 3 | version = "0.0.0" 4 | edition = "2021" 5 | publish = false 6 | authors = ["Miguel Piedrafita "] 7 | 8 | [package.metadata.cog] 9 | image = "blur-rs" 10 | 11 | [dependencies] 12 | serde = "1.0.163" 13 | image = "0.24.6" 14 | anyhow = "1.0.71" 15 | schemars = "0.8.12" 16 | cog-rust = { path = "../../lib" } 17 | tokio = { version = "1.28.2", features = ["full"] } 18 | 19 | # The image crate is extremely slow on debug builds, so we force a release build for better performance in development. 20 | # This does not affect the final release build. 21 | [profile.dev.package.image] 22 | opt-level = 3 23 | -------------------------------------------------------------------------------- /examples/blur/src/main.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use cog_rust::{Cog, Path}; 3 | use schemars::JsonSchema; 4 | 5 | #[derive(serde::Deserialize, JsonSchema)] 6 | struct ModelRequest { 7 | /// Input image 8 | image: Path, 9 | /// Blur radius (default: 5) 10 | blur: Option, 11 | } 12 | 13 | struct BlurModel {} 14 | 15 | impl Cog for BlurModel { 16 | type Request = ModelRequest; 17 | type Response = Path; 18 | 19 | async fn setup() -> Result { 20 | Ok(Self {}) 21 | } 22 | 23 | fn predict(&self, input: Self::Request) -> Result { 24 | let image = image::open(&input.image)?; 25 | image.blur(input.blur.unwrap_or(5.0)).save(&input.image)?; 26 | 27 | Ok(input.image) 28 | } 29 | } 30 | 31 | cog_rust::start!(BlurModel); 32 | -------------------------------------------------------------------------------- /examples/hello-world/Cargo.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Cargo. 2 | # It is not intended for manual editing. 3 | version = 3 4 | 5 | [[package]] 6 | name = "addr2line" 7 | version = "0.21.0" 8 | source = "registry+https://github.com/rust-lang/crates.io-index" 9 | checksum = "8a30b2e23b9e17a9f90641c7ab1549cd9b44f296d3ccbf309d2863cfe398a0cb" 10 | dependencies = [ 11 | "gimli", 12 | ] 13 | 14 | [[package]] 15 | name = "adler" 16 | version = "1.0.2" 17 | source = "registry+https://github.com/rust-lang/crates.io-index" 18 | checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" 19 | 20 | [[package]] 21 | name = "ahash" 22 | version = "0.8.3" 23 | source = "registry+https://github.com/rust-lang/crates.io-index" 24 | checksum = "2c99f64d1e06488f620f932677e24bc6e2897582980441ae90a671415bd7ec2f" 25 | dependencies = [ 26 | "cfg-if", 27 | "getrandom", 28 | "once_cell", 29 | "serde", 30 | "version_check", 31 | ] 32 | 33 | [[package]] 34 | name = "aho-corasick" 35 | version = "1.0.2" 36 | source = "registry+https://github.com/rust-lang/crates.io-index" 37 | checksum = "43f6cb1bf222025340178f382c426f13757b2960e89779dfcb319c32542a5a41" 38 | dependencies = [ 39 | "memchr", 40 | ] 41 | 42 | [[package]] 43 | name = "aide" 44 | version = "0.11.0" 45 | source = "registry+https://github.com/rust-lang/crates.io-index" 46 | checksum = "61ff6a9db74440d4847926f147f31973c60b8a852f838cb93652cbd0a28decc0" 47 | dependencies = [ 48 | "axum", 49 | "bytes", 50 | "cfg-if", 51 | "http", 52 | "indexmap", 53 | "schemars", 54 | "serde", 55 | "serde_json", 56 | "serde_qs", 57 | "thiserror", 58 | "tower-layer", 59 | "tower-service", 60 | "tracing", 61 | ] 62 | 63 | [[package]] 64 | name = "android-tzdata" 65 | version = "0.1.1" 66 | source = "registry+https://github.com/rust-lang/crates.io-index" 67 | checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" 68 | 69 | [[package]] 70 | name = "android_system_properties" 71 | version = "0.1.5" 72 | source = "registry+https://github.com/rust-lang/crates.io-index" 73 | checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" 74 | dependencies = [ 75 | "libc", 76 | ] 77 | 78 | [[package]] 79 | name = "anstream" 80 | version = "0.3.2" 81 | source = "registry+https://github.com/rust-lang/crates.io-index" 82 | checksum = "0ca84f3628370c59db74ee214b3263d58f9aadd9b4fe7e711fd87dc452b7f163" 83 | dependencies = [ 84 | "anstyle", 85 | "anstyle-parse", 86 | "anstyle-query", 87 | "anstyle-wincon", 88 | "colorchoice", 89 | "is-terminal", 90 | "utf8parse", 91 | ] 92 | 93 | [[package]] 94 | name = "anstyle" 95 | version = "1.0.0" 96 | source = "registry+https://github.com/rust-lang/crates.io-index" 97 | checksum = "41ed9a86bf92ae6580e0a31281f65a1b1d867c0cc68d5346e2ae128dddfa6a7d" 98 | 99 | [[package]] 100 | name = "anstyle-parse" 101 | version = "0.2.0" 102 | source = "registry+https://github.com/rust-lang/crates.io-index" 103 | checksum = "e765fd216e48e067936442276d1d57399e37bce53c264d6fefbe298080cb57ee" 104 | dependencies = [ 105 | "utf8parse", 106 | ] 107 | 108 | [[package]] 109 | name = "anstyle-query" 110 | version = "1.0.0" 111 | source = "registry+https://github.com/rust-lang/crates.io-index" 112 | checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b" 113 | dependencies = [ 114 | "windows-sys 0.48.0", 115 | ] 116 | 117 | [[package]] 118 | name = "anstyle-wincon" 119 | version = "1.0.1" 120 | source = "registry+https://github.com/rust-lang/crates.io-index" 121 | checksum = "180abfa45703aebe0093f79badacc01b8fd4ea2e35118747e5811127f926e188" 122 | dependencies = [ 123 | "anstyle", 124 | "windows-sys 0.48.0", 125 | ] 126 | 127 | [[package]] 128 | name = "anyhow" 129 | version = "1.0.71" 130 | source = "registry+https://github.com/rust-lang/crates.io-index" 131 | checksum = "9c7d0618f0e0b7e8ff11427422b64564d5fb0be1940354bfe2e0529b18a9d9b8" 132 | 133 | [[package]] 134 | name = "async-trait" 135 | version = "0.1.68" 136 | source = "registry+https://github.com/rust-lang/crates.io-index" 137 | checksum = "b9ccdd8f2a161be9bd5c023df56f1b2a0bd1d83872ae53b71a84a12c9bf6e842" 138 | dependencies = [ 139 | "proc-macro2", 140 | "quote", 141 | "syn 2.0.18", 142 | ] 143 | 144 | [[package]] 145 | name = "atomic_enum" 146 | version = "0.2.0" 147 | source = "registry+https://github.com/rust-lang/crates.io-index" 148 | checksum = "6227a8d6fdb862bcb100c4314d0d9579e5cd73fa6df31a2e6f6e1acd3c5f1207" 149 | dependencies = [ 150 | "proc-macro2", 151 | "quote", 152 | "syn 1.0.109", 153 | ] 154 | 155 | [[package]] 156 | name = "autocfg" 157 | version = "1.1.0" 158 | source = "registry+https://github.com/rust-lang/crates.io-index" 159 | checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" 160 | 161 | [[package]] 162 | name = "axum" 163 | version = "0.6.18" 164 | source = "registry+https://github.com/rust-lang/crates.io-index" 165 | checksum = "f8175979259124331c1d7bf6586ee7e0da434155e4b2d48ec2c8386281d8df39" 166 | dependencies = [ 167 | "async-trait", 168 | "axum-core", 169 | "bitflags", 170 | "bytes", 171 | "futures-util", 172 | "headers", 173 | "http", 174 | "http-body", 175 | "hyper", 176 | "itoa", 177 | "matchit", 178 | "memchr", 179 | "mime", 180 | "percent-encoding", 181 | "pin-project-lite", 182 | "rustversion", 183 | "serde", 184 | "serde_json", 185 | "serde_path_to_error", 186 | "serde_urlencoded", 187 | "sync_wrapper", 188 | "tokio", 189 | "tower", 190 | "tower-layer", 191 | "tower-service", 192 | ] 193 | 194 | [[package]] 195 | name = "axum-core" 196 | version = "0.3.4" 197 | source = "registry+https://github.com/rust-lang/crates.io-index" 198 | checksum = "759fa577a247914fd3f7f76d62972792636412fbfd634cd452f6a385a74d2d2c" 199 | dependencies = [ 200 | "async-trait", 201 | "bytes", 202 | "futures-util", 203 | "http", 204 | "http-body", 205 | "mime", 206 | "rustversion", 207 | "tower-layer", 208 | "tower-service", 209 | ] 210 | 211 | [[package]] 212 | name = "axum-jsonschema" 213 | version = "0.6.0" 214 | source = "registry+https://github.com/rust-lang/crates.io-index" 215 | checksum = "1c84a1b355067122ebbe8ae45ef7b290ff85c423342f0661116d1836a816ea70" 216 | dependencies = [ 217 | "aide", 218 | "async-trait", 219 | "axum", 220 | "http", 221 | "http-body", 222 | "itertools 0.10.5", 223 | "jsonschema", 224 | "schemars", 225 | "serde", 226 | "serde_json", 227 | "serde_path_to_error", 228 | "tracing", 229 | ] 230 | 231 | [[package]] 232 | name = "backtrace" 233 | version = "0.3.69" 234 | source = "registry+https://github.com/rust-lang/crates.io-index" 235 | checksum = "2089b7e3f35b9dd2d0ed921ead4f6d318c27680d4a5bd167b3ee120edb105837" 236 | dependencies = [ 237 | "addr2line", 238 | "cc", 239 | "cfg-if", 240 | "libc", 241 | "miniz_oxide", 242 | "object", 243 | "rustc-demangle", 244 | ] 245 | 246 | [[package]] 247 | name = "base64" 248 | version = "0.13.1" 249 | source = "registry+https://github.com/rust-lang/crates.io-index" 250 | checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" 251 | 252 | [[package]] 253 | name = "base64" 254 | version = "0.21.2" 255 | source = "registry+https://github.com/rust-lang/crates.io-index" 256 | checksum = "604178f6c5c21f02dc555784810edfb88d34ac2c73b2eae109655649ee73ce3d" 257 | 258 | [[package]] 259 | name = "bit-set" 260 | version = "0.5.3" 261 | source = "registry+https://github.com/rust-lang/crates.io-index" 262 | checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" 263 | dependencies = [ 264 | "bit-vec", 265 | ] 266 | 267 | [[package]] 268 | name = "bit-vec" 269 | version = "0.6.3" 270 | source = "registry+https://github.com/rust-lang/crates.io-index" 271 | checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" 272 | 273 | [[package]] 274 | name = "bitflags" 275 | version = "1.3.2" 276 | source = "registry+https://github.com/rust-lang/crates.io-index" 277 | checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" 278 | 279 | [[package]] 280 | name = "block-buffer" 281 | version = "0.10.4" 282 | source = "registry+https://github.com/rust-lang/crates.io-index" 283 | checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" 284 | dependencies = [ 285 | "generic-array", 286 | ] 287 | 288 | [[package]] 289 | name = "bumpalo" 290 | version = "3.13.0" 291 | source = "registry+https://github.com/rust-lang/crates.io-index" 292 | checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1" 293 | 294 | [[package]] 295 | name = "bytecount" 296 | version = "0.6.3" 297 | source = "registry+https://github.com/rust-lang/crates.io-index" 298 | checksum = "2c676a478f63e9fa2dd5368a42f28bba0d6c560b775f38583c8bbaa7fcd67c9c" 299 | 300 | [[package]] 301 | name = "bytes" 302 | version = "1.4.0" 303 | source = "registry+https://github.com/rust-lang/crates.io-index" 304 | checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" 305 | 306 | [[package]] 307 | name = "cc" 308 | version = "1.0.79" 309 | source = "registry+https://github.com/rust-lang/crates.io-index" 310 | checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f" 311 | 312 | [[package]] 313 | name = "cfg-if" 314 | version = "1.0.0" 315 | source = "registry+https://github.com/rust-lang/crates.io-index" 316 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 317 | 318 | [[package]] 319 | name = "chrono" 320 | version = "0.4.26" 321 | source = "registry+https://github.com/rust-lang/crates.io-index" 322 | checksum = "ec837a71355b28f6556dbd569b37b3f363091c0bd4b2e735674521b4c5fd9bc5" 323 | dependencies = [ 324 | "android-tzdata", 325 | "iana-time-zone", 326 | "js-sys", 327 | "num-traits", 328 | "serde", 329 | "time 0.1.45", 330 | "wasm-bindgen", 331 | "winapi", 332 | ] 333 | 334 | [[package]] 335 | name = "clap" 336 | version = "4.3.21" 337 | source = "registry+https://github.com/rust-lang/crates.io-index" 338 | checksum = "c27cdf28c0f604ba3f512b0c9a409f8de8513e4816705deb0498b627e7c3a3fd" 339 | dependencies = [ 340 | "clap_builder", 341 | "clap_derive", 342 | "once_cell", 343 | ] 344 | 345 | [[package]] 346 | name = "clap_builder" 347 | version = "4.3.21" 348 | source = "registry+https://github.com/rust-lang/crates.io-index" 349 | checksum = "08a9f1ab5e9f01a9b81f202e8562eb9a10de70abf9eaeac1be465c28b75aa4aa" 350 | dependencies = [ 351 | "anstream", 352 | "anstyle", 353 | "clap_lex", 354 | "strsim", 355 | ] 356 | 357 | [[package]] 358 | name = "clap_derive" 359 | version = "4.3.12" 360 | source = "registry+https://github.com/rust-lang/crates.io-index" 361 | checksum = "54a9bb5758fc5dfe728d1019941681eccaf0cf8a4189b692a0ee2f2ecf90a050" 362 | dependencies = [ 363 | "heck", 364 | "proc-macro2", 365 | "quote", 366 | "syn 2.0.18", 367 | ] 368 | 369 | [[package]] 370 | name = "clap_lex" 371 | version = "0.5.0" 372 | source = "registry+https://github.com/rust-lang/crates.io-index" 373 | checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b" 374 | 375 | [[package]] 376 | name = "cog-core" 377 | version = "0.1.1" 378 | dependencies = [ 379 | "anyhow", 380 | "chrono", 381 | "schemars", 382 | "serde", 383 | "serde_json", 384 | "thiserror", 385 | "tokio", 386 | "url", 387 | ] 388 | 389 | [[package]] 390 | name = "cog-hello-world" 391 | version = "0.0.0" 392 | dependencies = [ 393 | "anyhow", 394 | "cog-rust", 395 | "schemars", 396 | "serde", 397 | "tokio", 398 | ] 399 | 400 | [[package]] 401 | name = "cog-rust" 402 | version = "1.0.14" 403 | dependencies = [ 404 | "aide", 405 | "anyhow", 406 | "async-trait", 407 | "atomic_enum", 408 | "axum", 409 | "axum-jsonschema", 410 | "base64 0.21.2", 411 | "chrono", 412 | "clap", 413 | "cog-core", 414 | "flume", 415 | "indexmap", 416 | "itertools 0.11.0", 417 | "jsonschema", 418 | "lazy_static", 419 | "map-macro", 420 | "mime_guess", 421 | "percent-encoding", 422 | "reqwest", 423 | "schemars", 424 | "serde", 425 | "serde_json", 426 | "thiserror", 427 | "titlecase", 428 | "tokio", 429 | "tracing", 430 | "tracing-subscriber", 431 | "tree_magic_mini", 432 | "url", 433 | "uuid", 434 | ] 435 | 436 | [[package]] 437 | name = "colorchoice" 438 | version = "1.0.0" 439 | source = "registry+https://github.com/rust-lang/crates.io-index" 440 | checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" 441 | 442 | [[package]] 443 | name = "core-foundation" 444 | version = "0.9.3" 445 | source = "registry+https://github.com/rust-lang/crates.io-index" 446 | checksum = "194a7a9e6de53fa55116934067c844d9d749312f75c6f6d0980e8c252f8c2146" 447 | dependencies = [ 448 | "core-foundation-sys", 449 | "libc", 450 | ] 451 | 452 | [[package]] 453 | name = "core-foundation-sys" 454 | version = "0.8.4" 455 | source = "registry+https://github.com/rust-lang/crates.io-index" 456 | checksum = "e496a50fda8aacccc86d7529e2c1e0892dbd0f898a6b5645b5561b89c3210efa" 457 | 458 | [[package]] 459 | name = "cpufeatures" 460 | version = "0.2.7" 461 | source = "registry+https://github.com/rust-lang/crates.io-index" 462 | checksum = "3e4c1eaa2012c47becbbad2ab175484c2a84d1185b566fb2cc5b8707343dfe58" 463 | dependencies = [ 464 | "libc", 465 | ] 466 | 467 | [[package]] 468 | name = "crypto-common" 469 | version = "0.1.6" 470 | source = "registry+https://github.com/rust-lang/crates.io-index" 471 | checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" 472 | dependencies = [ 473 | "generic-array", 474 | "typenum", 475 | ] 476 | 477 | [[package]] 478 | name = "digest" 479 | version = "0.10.7" 480 | source = "registry+https://github.com/rust-lang/crates.io-index" 481 | checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" 482 | dependencies = [ 483 | "block-buffer", 484 | "crypto-common", 485 | ] 486 | 487 | [[package]] 488 | name = "dyn-clone" 489 | version = "1.0.11" 490 | source = "registry+https://github.com/rust-lang/crates.io-index" 491 | checksum = "68b0cf012f1230e43cd00ebb729c6bb58707ecfa8ad08b52ef3a4ccd2697fc30" 492 | 493 | [[package]] 494 | name = "either" 495 | version = "1.8.1" 496 | source = "registry+https://github.com/rust-lang/crates.io-index" 497 | checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91" 498 | 499 | [[package]] 500 | name = "encoding_rs" 501 | version = "0.8.32" 502 | source = "registry+https://github.com/rust-lang/crates.io-index" 503 | checksum = "071a31f4ee85403370b58aca746f01041ede6f0da2730960ad001edc2b71b394" 504 | dependencies = [ 505 | "cfg-if", 506 | ] 507 | 508 | [[package]] 509 | name = "errno" 510 | version = "0.3.1" 511 | source = "registry+https://github.com/rust-lang/crates.io-index" 512 | checksum = "4bcfec3a70f97c962c307b2d2c56e358cf1d00b558d74262b5f929ee8cc7e73a" 513 | dependencies = [ 514 | "errno-dragonfly", 515 | "libc", 516 | "windows-sys 0.48.0", 517 | ] 518 | 519 | [[package]] 520 | name = "errno-dragonfly" 521 | version = "0.1.2" 522 | source = "registry+https://github.com/rust-lang/crates.io-index" 523 | checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" 524 | dependencies = [ 525 | "cc", 526 | "libc", 527 | ] 528 | 529 | [[package]] 530 | name = "fancy-regex" 531 | version = "0.11.0" 532 | source = "registry+https://github.com/rust-lang/crates.io-index" 533 | checksum = "b95f7c0680e4142284cf8b22c14a476e87d61b004a3a0861872b32ef7ead40a2" 534 | dependencies = [ 535 | "bit-set", 536 | "regex", 537 | ] 538 | 539 | [[package]] 540 | name = "fastrand" 541 | version = "1.9.0" 542 | source = "registry+https://github.com/rust-lang/crates.io-index" 543 | checksum = "e51093e27b0797c359783294ca4f0a911c270184cb10f85783b118614a1501be" 544 | dependencies = [ 545 | "instant", 546 | ] 547 | 548 | [[package]] 549 | name = "fixedbitset" 550 | version = "0.4.2" 551 | source = "registry+https://github.com/rust-lang/crates.io-index" 552 | checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" 553 | 554 | [[package]] 555 | name = "flume" 556 | version = "0.10.14" 557 | source = "registry+https://github.com/rust-lang/crates.io-index" 558 | checksum = "1657b4441c3403d9f7b3409e47575237dac27b1b5726df654a6ecbf92f0f7577" 559 | dependencies = [ 560 | "futures-core", 561 | "futures-sink", 562 | "nanorand", 563 | "pin-project", 564 | "spin", 565 | ] 566 | 567 | [[package]] 568 | name = "fnv" 569 | version = "1.0.7" 570 | source = "registry+https://github.com/rust-lang/crates.io-index" 571 | checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" 572 | 573 | [[package]] 574 | name = "foreign-types" 575 | version = "0.3.2" 576 | source = "registry+https://github.com/rust-lang/crates.io-index" 577 | checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" 578 | dependencies = [ 579 | "foreign-types-shared", 580 | ] 581 | 582 | [[package]] 583 | name = "foreign-types-shared" 584 | version = "0.1.1" 585 | source = "registry+https://github.com/rust-lang/crates.io-index" 586 | checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" 587 | 588 | [[package]] 589 | name = "form_urlencoded" 590 | version = "1.2.0" 591 | source = "registry+https://github.com/rust-lang/crates.io-index" 592 | checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652" 593 | dependencies = [ 594 | "percent-encoding", 595 | ] 596 | 597 | [[package]] 598 | name = "fraction" 599 | version = "0.13.1" 600 | source = "registry+https://github.com/rust-lang/crates.io-index" 601 | checksum = "3027ae1df8d41b4bed2241c8fdad4acc1e7af60c8e17743534b545e77182d678" 602 | dependencies = [ 603 | "lazy_static", 604 | "num", 605 | ] 606 | 607 | [[package]] 608 | name = "futures" 609 | version = "0.3.28" 610 | source = "registry+https://github.com/rust-lang/crates.io-index" 611 | checksum = "23342abe12aba583913b2e62f22225ff9c950774065e4bfb61a19cd9770fec40" 612 | dependencies = [ 613 | "futures-channel", 614 | "futures-core", 615 | "futures-executor", 616 | "futures-io", 617 | "futures-sink", 618 | "futures-task", 619 | "futures-util", 620 | ] 621 | 622 | [[package]] 623 | name = "futures-channel" 624 | version = "0.3.28" 625 | source = "registry+https://github.com/rust-lang/crates.io-index" 626 | checksum = "955518d47e09b25bbebc7a18df10b81f0c766eaf4c4f1cccef2fca5f2a4fb5f2" 627 | dependencies = [ 628 | "futures-core", 629 | "futures-sink", 630 | ] 631 | 632 | [[package]] 633 | name = "futures-core" 634 | version = "0.3.28" 635 | source = "registry+https://github.com/rust-lang/crates.io-index" 636 | checksum = "4bca583b7e26f571124fe5b7561d49cb2868d79116cfa0eefce955557c6fee8c" 637 | 638 | [[package]] 639 | name = "futures-executor" 640 | version = "0.3.28" 641 | source = "registry+https://github.com/rust-lang/crates.io-index" 642 | checksum = "ccecee823288125bd88b4d7f565c9e58e41858e47ab72e8ea2d64e93624386e0" 643 | dependencies = [ 644 | "futures-core", 645 | "futures-task", 646 | "futures-util", 647 | ] 648 | 649 | [[package]] 650 | name = "futures-io" 651 | version = "0.3.28" 652 | source = "registry+https://github.com/rust-lang/crates.io-index" 653 | checksum = "4fff74096e71ed47f8e023204cfd0aa1289cd54ae5430a9523be060cdb849964" 654 | 655 | [[package]] 656 | name = "futures-macro" 657 | version = "0.3.28" 658 | source = "registry+https://github.com/rust-lang/crates.io-index" 659 | checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" 660 | dependencies = [ 661 | "proc-macro2", 662 | "quote", 663 | "syn 2.0.18", 664 | ] 665 | 666 | [[package]] 667 | name = "futures-sink" 668 | version = "0.3.28" 669 | source = "registry+https://github.com/rust-lang/crates.io-index" 670 | checksum = "f43be4fe21a13b9781a69afa4985b0f6ee0e1afab2c6f454a8cf30e2b2237b6e" 671 | 672 | [[package]] 673 | name = "futures-task" 674 | version = "0.3.28" 675 | source = "registry+https://github.com/rust-lang/crates.io-index" 676 | checksum = "76d3d132be6c0e6aa1534069c705a74a5997a356c0dc2f86a47765e5617c5b65" 677 | 678 | [[package]] 679 | name = "futures-util" 680 | version = "0.3.28" 681 | source = "registry+https://github.com/rust-lang/crates.io-index" 682 | checksum = "26b01e40b772d54cf6c6d721c1d1abd0647a0106a12ecaa1c186273392a69533" 683 | dependencies = [ 684 | "futures-channel", 685 | "futures-core", 686 | "futures-io", 687 | "futures-macro", 688 | "futures-sink", 689 | "futures-task", 690 | "memchr", 691 | "pin-project-lite", 692 | "pin-utils", 693 | "slab", 694 | ] 695 | 696 | [[package]] 697 | name = "generic-array" 698 | version = "0.14.7" 699 | source = "registry+https://github.com/rust-lang/crates.io-index" 700 | checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" 701 | dependencies = [ 702 | "typenum", 703 | "version_check", 704 | ] 705 | 706 | [[package]] 707 | name = "getrandom" 708 | version = "0.2.10" 709 | source = "registry+https://github.com/rust-lang/crates.io-index" 710 | checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" 711 | dependencies = [ 712 | "cfg-if", 713 | "js-sys", 714 | "libc", 715 | "wasi 0.11.0+wasi-snapshot-preview1", 716 | "wasm-bindgen", 717 | ] 718 | 719 | [[package]] 720 | name = "gimli" 721 | version = "0.28.1" 722 | source = "registry+https://github.com/rust-lang/crates.io-index" 723 | checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253" 724 | 725 | [[package]] 726 | name = "h2" 727 | version = "0.3.19" 728 | source = "registry+https://github.com/rust-lang/crates.io-index" 729 | checksum = "d357c7ae988e7d2182f7d7871d0b963962420b0678b0997ce7de72001aeab782" 730 | dependencies = [ 731 | "bytes", 732 | "fnv", 733 | "futures-core", 734 | "futures-sink", 735 | "futures-util", 736 | "http", 737 | "indexmap", 738 | "slab", 739 | "tokio", 740 | "tokio-util", 741 | "tracing", 742 | ] 743 | 744 | [[package]] 745 | name = "hashbrown" 746 | version = "0.12.3" 747 | source = "registry+https://github.com/rust-lang/crates.io-index" 748 | checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" 749 | 750 | [[package]] 751 | name = "headers" 752 | version = "0.3.8" 753 | source = "registry+https://github.com/rust-lang/crates.io-index" 754 | checksum = "f3e372db8e5c0d213e0cd0b9be18be2aca3d44cf2fe30a9d46a65581cd454584" 755 | dependencies = [ 756 | "base64 0.13.1", 757 | "bitflags", 758 | "bytes", 759 | "headers-core", 760 | "http", 761 | "httpdate", 762 | "mime", 763 | "sha1", 764 | ] 765 | 766 | [[package]] 767 | name = "headers-core" 768 | version = "0.2.0" 769 | source = "registry+https://github.com/rust-lang/crates.io-index" 770 | checksum = "e7f66481bfee273957b1f20485a4ff3362987f85b2c236580d81b4eb7a326429" 771 | dependencies = [ 772 | "http", 773 | ] 774 | 775 | [[package]] 776 | name = "heck" 777 | version = "0.4.1" 778 | source = "registry+https://github.com/rust-lang/crates.io-index" 779 | checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" 780 | 781 | [[package]] 782 | name = "hermit-abi" 783 | version = "0.2.6" 784 | source = "registry+https://github.com/rust-lang/crates.io-index" 785 | checksum = "ee512640fe35acbfb4bb779db6f0d80704c2cacfa2e39b601ef3e3f47d1ae4c7" 786 | dependencies = [ 787 | "libc", 788 | ] 789 | 790 | [[package]] 791 | name = "hermit-abi" 792 | version = "0.3.1" 793 | source = "registry+https://github.com/rust-lang/crates.io-index" 794 | checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286" 795 | 796 | [[package]] 797 | name = "http" 798 | version = "0.2.9" 799 | source = "registry+https://github.com/rust-lang/crates.io-index" 800 | checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482" 801 | dependencies = [ 802 | "bytes", 803 | "fnv", 804 | "itoa", 805 | ] 806 | 807 | [[package]] 808 | name = "http-body" 809 | version = "0.4.5" 810 | source = "registry+https://github.com/rust-lang/crates.io-index" 811 | checksum = "d5f38f16d184e36f2408a55281cd658ecbd3ca05cce6d6510a176eca393e26d1" 812 | dependencies = [ 813 | "bytes", 814 | "http", 815 | "pin-project-lite", 816 | ] 817 | 818 | [[package]] 819 | name = "httparse" 820 | version = "1.8.0" 821 | source = "registry+https://github.com/rust-lang/crates.io-index" 822 | checksum = "d897f394bad6a705d5f4104762e116a75639e470d80901eed05a860a95cb1904" 823 | 824 | [[package]] 825 | name = "httpdate" 826 | version = "1.0.2" 827 | source = "registry+https://github.com/rust-lang/crates.io-index" 828 | checksum = "c4a1e36c821dbe04574f602848a19f742f4fb3c98d40449f11bcad18d6b17421" 829 | 830 | [[package]] 831 | name = "hyper" 832 | version = "0.14.26" 833 | source = "registry+https://github.com/rust-lang/crates.io-index" 834 | checksum = "ab302d72a6f11a3b910431ff93aae7e773078c769f0a3ef15fb9ec692ed147d4" 835 | dependencies = [ 836 | "bytes", 837 | "futures-channel", 838 | "futures-core", 839 | "futures-util", 840 | "h2", 841 | "http", 842 | "http-body", 843 | "httparse", 844 | "httpdate", 845 | "itoa", 846 | "pin-project-lite", 847 | "socket2 0.4.9", 848 | "tokio", 849 | "tower-service", 850 | "tracing", 851 | "want", 852 | ] 853 | 854 | [[package]] 855 | name = "hyper-tls" 856 | version = "0.5.0" 857 | source = "registry+https://github.com/rust-lang/crates.io-index" 858 | checksum = "d6183ddfa99b85da61a140bea0efc93fdf56ceaa041b37d553518030827f9905" 859 | dependencies = [ 860 | "bytes", 861 | "hyper", 862 | "native-tls", 863 | "tokio", 864 | "tokio-native-tls", 865 | ] 866 | 867 | [[package]] 868 | name = "iana-time-zone" 869 | version = "0.1.57" 870 | source = "registry+https://github.com/rust-lang/crates.io-index" 871 | checksum = "2fad5b825842d2b38bd206f3e81d6957625fd7f0a361e345c30e01a0ae2dd613" 872 | dependencies = [ 873 | "android_system_properties", 874 | "core-foundation-sys", 875 | "iana-time-zone-haiku", 876 | "js-sys", 877 | "wasm-bindgen", 878 | "windows", 879 | ] 880 | 881 | [[package]] 882 | name = "iana-time-zone-haiku" 883 | version = "0.1.2" 884 | source = "registry+https://github.com/rust-lang/crates.io-index" 885 | checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" 886 | dependencies = [ 887 | "cc", 888 | ] 889 | 890 | [[package]] 891 | name = "idna" 892 | version = "0.4.0" 893 | source = "registry+https://github.com/rust-lang/crates.io-index" 894 | checksum = "7d20d6b07bfbc108882d88ed8e37d39636dcc260e15e30c45e6ba089610b917c" 895 | dependencies = [ 896 | "unicode-bidi", 897 | "unicode-normalization", 898 | ] 899 | 900 | [[package]] 901 | name = "indexmap" 902 | version = "1.9.3" 903 | source = "registry+https://github.com/rust-lang/crates.io-index" 904 | checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" 905 | dependencies = [ 906 | "autocfg", 907 | "hashbrown", 908 | "serde", 909 | ] 910 | 911 | [[package]] 912 | name = "instant" 913 | version = "0.1.12" 914 | source = "registry+https://github.com/rust-lang/crates.io-index" 915 | checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" 916 | dependencies = [ 917 | "cfg-if", 918 | ] 919 | 920 | [[package]] 921 | name = "io-lifetimes" 922 | version = "1.0.11" 923 | source = "registry+https://github.com/rust-lang/crates.io-index" 924 | checksum = "eae7b9aee968036d54dce06cebaefd919e4472e753296daccd6d344e3e2df0c2" 925 | dependencies = [ 926 | "hermit-abi 0.3.1", 927 | "libc", 928 | "windows-sys 0.48.0", 929 | ] 930 | 931 | [[package]] 932 | name = "ipnet" 933 | version = "2.7.2" 934 | source = "registry+https://github.com/rust-lang/crates.io-index" 935 | checksum = "12b6ee2129af8d4fb011108c73d99a1b83a85977f23b82460c0ae2e25bb4b57f" 936 | 937 | [[package]] 938 | name = "is-terminal" 939 | version = "0.4.7" 940 | source = "registry+https://github.com/rust-lang/crates.io-index" 941 | checksum = "adcf93614601c8129ddf72e2d5633df827ba6551541c6d8c59520a371475be1f" 942 | dependencies = [ 943 | "hermit-abi 0.3.1", 944 | "io-lifetimes", 945 | "rustix", 946 | "windows-sys 0.48.0", 947 | ] 948 | 949 | [[package]] 950 | name = "iso8601" 951 | version = "0.6.1" 952 | source = "registry+https://github.com/rust-lang/crates.io-index" 953 | checksum = "924e5d73ea28f59011fec52a0d12185d496a9b075d360657aed2a5707f701153" 954 | dependencies = [ 955 | "nom", 956 | ] 957 | 958 | [[package]] 959 | name = "itertools" 960 | version = "0.10.5" 961 | source = "registry+https://github.com/rust-lang/crates.io-index" 962 | checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" 963 | dependencies = [ 964 | "either", 965 | ] 966 | 967 | [[package]] 968 | name = "itertools" 969 | version = "0.11.0" 970 | source = "registry+https://github.com/rust-lang/crates.io-index" 971 | checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" 972 | dependencies = [ 973 | "either", 974 | ] 975 | 976 | [[package]] 977 | name = "itoa" 978 | version = "1.0.6" 979 | source = "registry+https://github.com/rust-lang/crates.io-index" 980 | checksum = "453ad9f582a441959e5f0d088b02ce04cfe8d51a8eaf077f12ac6d3e94164ca6" 981 | 982 | [[package]] 983 | name = "joinery" 984 | version = "2.1.0" 985 | source = "registry+https://github.com/rust-lang/crates.io-index" 986 | checksum = "72167d68f5fce3b8655487b8038691a3c9984ee769590f93f2a631f4ad64e4f5" 987 | 988 | [[package]] 989 | name = "js-sys" 990 | version = "0.3.63" 991 | source = "registry+https://github.com/rust-lang/crates.io-index" 992 | checksum = "2f37a4a5928311ac501dee68b3c7613a1037d0edb30c8e5427bd832d55d1b790" 993 | dependencies = [ 994 | "wasm-bindgen", 995 | ] 996 | 997 | [[package]] 998 | name = "jsonschema" 999 | version = "0.17.0" 1000 | source = "registry+https://github.com/rust-lang/crates.io-index" 1001 | checksum = "e48354c4c4f088714424ddf090de1ff84acc82b2f08c192d46d226ae2529a465" 1002 | dependencies = [ 1003 | "ahash", 1004 | "anyhow", 1005 | "base64 0.21.2", 1006 | "bytecount", 1007 | "fancy-regex", 1008 | "fraction", 1009 | "getrandom", 1010 | "iso8601", 1011 | "itoa", 1012 | "memchr", 1013 | "num-cmp", 1014 | "once_cell", 1015 | "parking_lot", 1016 | "percent-encoding", 1017 | "regex", 1018 | "serde", 1019 | "serde_json", 1020 | "time 0.3.22", 1021 | "url", 1022 | "uuid", 1023 | ] 1024 | 1025 | [[package]] 1026 | name = "lazy_static" 1027 | version = "1.4.0" 1028 | source = "registry+https://github.com/rust-lang/crates.io-index" 1029 | checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" 1030 | 1031 | [[package]] 1032 | name = "libc" 1033 | version = "0.2.151" 1034 | source = "registry+https://github.com/rust-lang/crates.io-index" 1035 | checksum = "302d7ab3130588088d277783b1e2d2e10c9e9e4a16dd9050e6ec93fb3e7048f4" 1036 | 1037 | [[package]] 1038 | name = "linux-raw-sys" 1039 | version = "0.3.8" 1040 | source = "registry+https://github.com/rust-lang/crates.io-index" 1041 | checksum = "ef53942eb7bf7ff43a617b3e2c1c4a5ecf5944a7c1bc12d7ee39bbb15e5c1519" 1042 | 1043 | [[package]] 1044 | name = "lock_api" 1045 | version = "0.4.10" 1046 | source = "registry+https://github.com/rust-lang/crates.io-index" 1047 | checksum = "c1cc9717a20b1bb222f333e6a92fd32f7d8a18ddc5a3191a11af45dcbf4dcd16" 1048 | dependencies = [ 1049 | "autocfg", 1050 | "scopeguard", 1051 | ] 1052 | 1053 | [[package]] 1054 | name = "log" 1055 | version = "0.4.18" 1056 | source = "registry+https://github.com/rust-lang/crates.io-index" 1057 | checksum = "518ef76f2f87365916b142844c16d8fefd85039bc5699050210a7778ee1cd1de" 1058 | 1059 | [[package]] 1060 | name = "map-macro" 1061 | version = "0.2.6" 1062 | source = "registry+https://github.com/rust-lang/crates.io-index" 1063 | checksum = "7c2efbd1385acc8dad8a2e56558e58d949d777741fe110f2ddf3472671dbe3e6" 1064 | 1065 | [[package]] 1066 | name = "matchers" 1067 | version = "0.1.0" 1068 | source = "registry+https://github.com/rust-lang/crates.io-index" 1069 | checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" 1070 | dependencies = [ 1071 | "regex-automata", 1072 | ] 1073 | 1074 | [[package]] 1075 | name = "matchit" 1076 | version = "0.7.0" 1077 | source = "registry+https://github.com/rust-lang/crates.io-index" 1078 | checksum = "b87248edafb776e59e6ee64a79086f65890d3510f2c656c000bf2a7e8a0aea40" 1079 | 1080 | [[package]] 1081 | name = "memchr" 1082 | version = "2.5.0" 1083 | source = "registry+https://github.com/rust-lang/crates.io-index" 1084 | checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" 1085 | 1086 | [[package]] 1087 | name = "mime" 1088 | version = "0.3.17" 1089 | source = "registry+https://github.com/rust-lang/crates.io-index" 1090 | checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" 1091 | 1092 | [[package]] 1093 | name = "mime_guess" 1094 | version = "2.0.4" 1095 | source = "registry+https://github.com/rust-lang/crates.io-index" 1096 | checksum = "4192263c238a5f0d0c6bfd21f336a313a4ce1c450542449ca191bb657b4642ef" 1097 | dependencies = [ 1098 | "mime", 1099 | "unicase", 1100 | ] 1101 | 1102 | [[package]] 1103 | name = "minimal-lexical" 1104 | version = "0.2.1" 1105 | source = "registry+https://github.com/rust-lang/crates.io-index" 1106 | checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" 1107 | 1108 | [[package]] 1109 | name = "miniz_oxide" 1110 | version = "0.7.1" 1111 | source = "registry+https://github.com/rust-lang/crates.io-index" 1112 | checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" 1113 | dependencies = [ 1114 | "adler", 1115 | ] 1116 | 1117 | [[package]] 1118 | name = "mio" 1119 | version = "0.8.10" 1120 | source = "registry+https://github.com/rust-lang/crates.io-index" 1121 | checksum = "8f3d0b296e374a4e6f3c7b0a1f5a51d748a0d34c85e7dc48fc3fa9a87657fe09" 1122 | dependencies = [ 1123 | "libc", 1124 | "wasi 0.11.0+wasi-snapshot-preview1", 1125 | "windows-sys 0.48.0", 1126 | ] 1127 | 1128 | [[package]] 1129 | name = "nanorand" 1130 | version = "0.7.0" 1131 | source = "registry+https://github.com/rust-lang/crates.io-index" 1132 | checksum = "6a51313c5820b0b02bd422f4b44776fbf47961755c74ce64afc73bfad10226c3" 1133 | dependencies = [ 1134 | "getrandom", 1135 | ] 1136 | 1137 | [[package]] 1138 | name = "native-tls" 1139 | version = "0.2.11" 1140 | source = "registry+https://github.com/rust-lang/crates.io-index" 1141 | checksum = "07226173c32f2926027b63cce4bcd8076c3552846cbe7925f3aaffeac0a3b92e" 1142 | dependencies = [ 1143 | "lazy_static", 1144 | "libc", 1145 | "log", 1146 | "openssl", 1147 | "openssl-probe", 1148 | "openssl-sys", 1149 | "schannel", 1150 | "security-framework", 1151 | "security-framework-sys", 1152 | "tempfile", 1153 | ] 1154 | 1155 | [[package]] 1156 | name = "nom" 1157 | version = "7.1.3" 1158 | source = "registry+https://github.com/rust-lang/crates.io-index" 1159 | checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" 1160 | dependencies = [ 1161 | "memchr", 1162 | "minimal-lexical", 1163 | ] 1164 | 1165 | [[package]] 1166 | name = "nu-ansi-term" 1167 | version = "0.46.0" 1168 | source = "registry+https://github.com/rust-lang/crates.io-index" 1169 | checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" 1170 | dependencies = [ 1171 | "overload", 1172 | "winapi", 1173 | ] 1174 | 1175 | [[package]] 1176 | name = "num" 1177 | version = "0.4.0" 1178 | source = "registry+https://github.com/rust-lang/crates.io-index" 1179 | checksum = "43db66d1170d347f9a065114077f7dccb00c1b9478c89384490a3425279a4606" 1180 | dependencies = [ 1181 | "num-bigint", 1182 | "num-complex", 1183 | "num-integer", 1184 | "num-iter", 1185 | "num-rational", 1186 | "num-traits", 1187 | ] 1188 | 1189 | [[package]] 1190 | name = "num-bigint" 1191 | version = "0.4.3" 1192 | source = "registry+https://github.com/rust-lang/crates.io-index" 1193 | checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f" 1194 | dependencies = [ 1195 | "autocfg", 1196 | "num-integer", 1197 | "num-traits", 1198 | ] 1199 | 1200 | [[package]] 1201 | name = "num-cmp" 1202 | version = "0.1.0" 1203 | source = "registry+https://github.com/rust-lang/crates.io-index" 1204 | checksum = "63335b2e2c34fae2fb0aa2cecfd9f0832a1e24b3b32ecec612c3426d46dc8aaa" 1205 | 1206 | [[package]] 1207 | name = "num-complex" 1208 | version = "0.4.3" 1209 | source = "registry+https://github.com/rust-lang/crates.io-index" 1210 | checksum = "02e0d21255c828d6f128a1e41534206671e8c3ea0c62f32291e808dc82cff17d" 1211 | dependencies = [ 1212 | "num-traits", 1213 | ] 1214 | 1215 | [[package]] 1216 | name = "num-integer" 1217 | version = "0.1.45" 1218 | source = "registry+https://github.com/rust-lang/crates.io-index" 1219 | checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" 1220 | dependencies = [ 1221 | "autocfg", 1222 | "num-traits", 1223 | ] 1224 | 1225 | [[package]] 1226 | name = "num-iter" 1227 | version = "0.1.43" 1228 | source = "registry+https://github.com/rust-lang/crates.io-index" 1229 | checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" 1230 | dependencies = [ 1231 | "autocfg", 1232 | "num-integer", 1233 | "num-traits", 1234 | ] 1235 | 1236 | [[package]] 1237 | name = "num-rational" 1238 | version = "0.4.1" 1239 | source = "registry+https://github.com/rust-lang/crates.io-index" 1240 | checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0" 1241 | dependencies = [ 1242 | "autocfg", 1243 | "num-bigint", 1244 | "num-integer", 1245 | "num-traits", 1246 | ] 1247 | 1248 | [[package]] 1249 | name = "num-traits" 1250 | version = "0.2.15" 1251 | source = "registry+https://github.com/rust-lang/crates.io-index" 1252 | checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" 1253 | dependencies = [ 1254 | "autocfg", 1255 | ] 1256 | 1257 | [[package]] 1258 | name = "num_cpus" 1259 | version = "1.15.0" 1260 | source = "registry+https://github.com/rust-lang/crates.io-index" 1261 | checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b" 1262 | dependencies = [ 1263 | "hermit-abi 0.2.6", 1264 | "libc", 1265 | ] 1266 | 1267 | [[package]] 1268 | name = "object" 1269 | version = "0.32.1" 1270 | source = "registry+https://github.com/rust-lang/crates.io-index" 1271 | checksum = "9cf5f9dd3933bd50a9e1f149ec995f39ae2c496d31fd772c1fd45ebc27e902b0" 1272 | dependencies = [ 1273 | "memchr", 1274 | ] 1275 | 1276 | [[package]] 1277 | name = "once_cell" 1278 | version = "1.18.0" 1279 | source = "registry+https://github.com/rust-lang/crates.io-index" 1280 | checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" 1281 | 1282 | [[package]] 1283 | name = "openssl" 1284 | version = "0.10.54" 1285 | source = "registry+https://github.com/rust-lang/crates.io-index" 1286 | checksum = "69b3f656a17a6cbc115b5c7a40c616947d213ba182135b014d6051b73ab6f019" 1287 | dependencies = [ 1288 | "bitflags", 1289 | "cfg-if", 1290 | "foreign-types", 1291 | "libc", 1292 | "once_cell", 1293 | "openssl-macros", 1294 | "openssl-sys", 1295 | ] 1296 | 1297 | [[package]] 1298 | name = "openssl-macros" 1299 | version = "0.1.1" 1300 | source = "registry+https://github.com/rust-lang/crates.io-index" 1301 | checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" 1302 | dependencies = [ 1303 | "proc-macro2", 1304 | "quote", 1305 | "syn 2.0.18", 1306 | ] 1307 | 1308 | [[package]] 1309 | name = "openssl-probe" 1310 | version = "0.1.5" 1311 | source = "registry+https://github.com/rust-lang/crates.io-index" 1312 | checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" 1313 | 1314 | [[package]] 1315 | name = "openssl-sys" 1316 | version = "0.9.88" 1317 | source = "registry+https://github.com/rust-lang/crates.io-index" 1318 | checksum = "c2ce0f250f34a308dcfdbb351f511359857d4ed2134ba715a4eadd46e1ffd617" 1319 | dependencies = [ 1320 | "cc", 1321 | "libc", 1322 | "pkg-config", 1323 | "vcpkg", 1324 | ] 1325 | 1326 | [[package]] 1327 | name = "overload" 1328 | version = "0.1.1" 1329 | source = "registry+https://github.com/rust-lang/crates.io-index" 1330 | checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" 1331 | 1332 | [[package]] 1333 | name = "parking_lot" 1334 | version = "0.12.1" 1335 | source = "registry+https://github.com/rust-lang/crates.io-index" 1336 | checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" 1337 | dependencies = [ 1338 | "lock_api", 1339 | "parking_lot_core", 1340 | ] 1341 | 1342 | [[package]] 1343 | name = "parking_lot_core" 1344 | version = "0.9.8" 1345 | source = "registry+https://github.com/rust-lang/crates.io-index" 1346 | checksum = "93f00c865fe7cabf650081affecd3871070f26767e7b2070a3ffae14c654b447" 1347 | dependencies = [ 1348 | "cfg-if", 1349 | "libc", 1350 | "redox_syscall", 1351 | "smallvec", 1352 | "windows-targets", 1353 | ] 1354 | 1355 | [[package]] 1356 | name = "percent-encoding" 1357 | version = "2.3.0" 1358 | source = "registry+https://github.com/rust-lang/crates.io-index" 1359 | checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" 1360 | 1361 | [[package]] 1362 | name = "petgraph" 1363 | version = "0.6.3" 1364 | source = "registry+https://github.com/rust-lang/crates.io-index" 1365 | checksum = "4dd7d28ee937e54fe3080c91faa1c3a46c06de6252988a7f4592ba2310ef22a4" 1366 | dependencies = [ 1367 | "fixedbitset", 1368 | "indexmap", 1369 | ] 1370 | 1371 | [[package]] 1372 | name = "pin-project" 1373 | version = "1.1.0" 1374 | source = "registry+https://github.com/rust-lang/crates.io-index" 1375 | checksum = "c95a7476719eab1e366eaf73d0260af3021184f18177925b07f54b30089ceead" 1376 | dependencies = [ 1377 | "pin-project-internal", 1378 | ] 1379 | 1380 | [[package]] 1381 | name = "pin-project-internal" 1382 | version = "1.1.0" 1383 | source = "registry+https://github.com/rust-lang/crates.io-index" 1384 | checksum = "39407670928234ebc5e6e580247dd567ad73a3578460c5990f9503df207e8f07" 1385 | dependencies = [ 1386 | "proc-macro2", 1387 | "quote", 1388 | "syn 2.0.18", 1389 | ] 1390 | 1391 | [[package]] 1392 | name = "pin-project-lite" 1393 | version = "0.2.13" 1394 | source = "registry+https://github.com/rust-lang/crates.io-index" 1395 | checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" 1396 | 1397 | [[package]] 1398 | name = "pin-utils" 1399 | version = "0.1.0" 1400 | source = "registry+https://github.com/rust-lang/crates.io-index" 1401 | checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" 1402 | 1403 | [[package]] 1404 | name = "pkg-config" 1405 | version = "0.3.27" 1406 | source = "registry+https://github.com/rust-lang/crates.io-index" 1407 | checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" 1408 | 1409 | [[package]] 1410 | name = "proc-macro2" 1411 | version = "1.0.60" 1412 | source = "registry+https://github.com/rust-lang/crates.io-index" 1413 | checksum = "dec2b086b7a862cf4de201096214fa870344cf922b2b30c167badb3af3195406" 1414 | dependencies = [ 1415 | "unicode-ident", 1416 | ] 1417 | 1418 | [[package]] 1419 | name = "quote" 1420 | version = "1.0.28" 1421 | source = "registry+https://github.com/rust-lang/crates.io-index" 1422 | checksum = "1b9ab9c7eadfd8df19006f1cf1a4aed13540ed5cbc047010ece5826e10825488" 1423 | dependencies = [ 1424 | "proc-macro2", 1425 | ] 1426 | 1427 | [[package]] 1428 | name = "redox_syscall" 1429 | version = "0.3.5" 1430 | source = "registry+https://github.com/rust-lang/crates.io-index" 1431 | checksum = "567664f262709473930a4bf9e51bf2ebf3348f2e748ccc50dea20646858f8f29" 1432 | dependencies = [ 1433 | "bitflags", 1434 | ] 1435 | 1436 | [[package]] 1437 | name = "regex" 1438 | version = "1.8.4" 1439 | source = "registry+https://github.com/rust-lang/crates.io-index" 1440 | checksum = "d0ab3ca65655bb1e41f2a8c8cd662eb4fb035e67c3f78da1d61dffe89d07300f" 1441 | dependencies = [ 1442 | "aho-corasick", 1443 | "memchr", 1444 | "regex-syntax 0.7.2", 1445 | ] 1446 | 1447 | [[package]] 1448 | name = "regex-automata" 1449 | version = "0.1.10" 1450 | source = "registry+https://github.com/rust-lang/crates.io-index" 1451 | checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" 1452 | dependencies = [ 1453 | "regex-syntax 0.6.29", 1454 | ] 1455 | 1456 | [[package]] 1457 | name = "regex-syntax" 1458 | version = "0.6.29" 1459 | source = "registry+https://github.com/rust-lang/crates.io-index" 1460 | checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" 1461 | 1462 | [[package]] 1463 | name = "regex-syntax" 1464 | version = "0.7.2" 1465 | source = "registry+https://github.com/rust-lang/crates.io-index" 1466 | checksum = "436b050e76ed2903236f032a59761c1eb99e1b0aead2c257922771dab1fc8c78" 1467 | 1468 | [[package]] 1469 | name = "reqwest" 1470 | version = "0.11.18" 1471 | source = "registry+https://github.com/rust-lang/crates.io-index" 1472 | checksum = "cde824a14b7c14f85caff81225f411faacc04a2013f41670f41443742b1c1c55" 1473 | dependencies = [ 1474 | "base64 0.21.2", 1475 | "bytes", 1476 | "encoding_rs", 1477 | "futures-core", 1478 | "futures-util", 1479 | "h2", 1480 | "http", 1481 | "http-body", 1482 | "hyper", 1483 | "hyper-tls", 1484 | "ipnet", 1485 | "js-sys", 1486 | "log", 1487 | "mime", 1488 | "native-tls", 1489 | "once_cell", 1490 | "percent-encoding", 1491 | "pin-project-lite", 1492 | "serde", 1493 | "serde_json", 1494 | "serde_urlencoded", 1495 | "tokio", 1496 | "tokio-native-tls", 1497 | "tower-service", 1498 | "url", 1499 | "wasm-bindgen", 1500 | "wasm-bindgen-futures", 1501 | "web-sys", 1502 | "winreg", 1503 | ] 1504 | 1505 | [[package]] 1506 | name = "rustc-demangle" 1507 | version = "0.1.23" 1508 | source = "registry+https://github.com/rust-lang/crates.io-index" 1509 | checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" 1510 | 1511 | [[package]] 1512 | name = "rustix" 1513 | version = "0.37.19" 1514 | source = "registry+https://github.com/rust-lang/crates.io-index" 1515 | checksum = "acf8729d8542766f1b2cf77eb034d52f40d375bb8b615d0b147089946e16613d" 1516 | dependencies = [ 1517 | "bitflags", 1518 | "errno", 1519 | "io-lifetimes", 1520 | "libc", 1521 | "linux-raw-sys", 1522 | "windows-sys 0.48.0", 1523 | ] 1524 | 1525 | [[package]] 1526 | name = "rustversion" 1527 | version = "1.0.12" 1528 | source = "registry+https://github.com/rust-lang/crates.io-index" 1529 | checksum = "4f3208ce4d8448b3f3e7d168a73f5e0c43a61e32930de3bceeccedb388b6bf06" 1530 | 1531 | [[package]] 1532 | name = "ryu" 1533 | version = "1.0.13" 1534 | source = "registry+https://github.com/rust-lang/crates.io-index" 1535 | checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041" 1536 | 1537 | [[package]] 1538 | name = "schannel" 1539 | version = "0.1.21" 1540 | source = "registry+https://github.com/rust-lang/crates.io-index" 1541 | checksum = "713cfb06c7059f3588fb8044c0fad1d09e3c01d225e25b9220dbfdcf16dbb1b3" 1542 | dependencies = [ 1543 | "windows-sys 0.42.0", 1544 | ] 1545 | 1546 | [[package]] 1547 | name = "schemars" 1548 | version = "0.8.12" 1549 | source = "registry+https://github.com/rust-lang/crates.io-index" 1550 | checksum = "02c613288622e5f0c3fdc5dbd4db1c5fbe752746b1d1a56a0630b78fd00de44f" 1551 | dependencies = [ 1552 | "chrono", 1553 | "dyn-clone", 1554 | "indexmap", 1555 | "schemars_derive", 1556 | "serde", 1557 | "serde_json", 1558 | "url", 1559 | ] 1560 | 1561 | [[package]] 1562 | name = "schemars_derive" 1563 | version = "0.8.12" 1564 | source = "registry+https://github.com/rust-lang/crates.io-index" 1565 | checksum = "109da1e6b197438deb6db99952990c7f959572794b80ff93707d55a232545e7c" 1566 | dependencies = [ 1567 | "proc-macro2", 1568 | "quote", 1569 | "serde_derive_internals", 1570 | "syn 1.0.109", 1571 | ] 1572 | 1573 | [[package]] 1574 | name = "scopeguard" 1575 | version = "1.1.0" 1576 | source = "registry+https://github.com/rust-lang/crates.io-index" 1577 | checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" 1578 | 1579 | [[package]] 1580 | name = "security-framework" 1581 | version = "2.9.1" 1582 | source = "registry+https://github.com/rust-lang/crates.io-index" 1583 | checksum = "1fc758eb7bffce5b308734e9b0c1468893cae9ff70ebf13e7090be8dcbcc83a8" 1584 | dependencies = [ 1585 | "bitflags", 1586 | "core-foundation", 1587 | "core-foundation-sys", 1588 | "libc", 1589 | "security-framework-sys", 1590 | ] 1591 | 1592 | [[package]] 1593 | name = "security-framework-sys" 1594 | version = "2.9.0" 1595 | source = "registry+https://github.com/rust-lang/crates.io-index" 1596 | checksum = "f51d0c0d83bec45f16480d0ce0058397a69e48fcdc52d1dc8855fb68acbd31a7" 1597 | dependencies = [ 1598 | "core-foundation-sys", 1599 | "libc", 1600 | ] 1601 | 1602 | [[package]] 1603 | name = "serde" 1604 | version = "1.0.164" 1605 | source = "registry+https://github.com/rust-lang/crates.io-index" 1606 | checksum = "9e8c8cf938e98f769bc164923b06dce91cea1751522f46f8466461af04c9027d" 1607 | dependencies = [ 1608 | "serde_derive", 1609 | ] 1610 | 1611 | [[package]] 1612 | name = "serde_derive" 1613 | version = "1.0.164" 1614 | source = "registry+https://github.com/rust-lang/crates.io-index" 1615 | checksum = "d9735b638ccc51c28bf6914d90a2e9725b377144fc612c49a611fddd1b631d68" 1616 | dependencies = [ 1617 | "proc-macro2", 1618 | "quote", 1619 | "syn 2.0.18", 1620 | ] 1621 | 1622 | [[package]] 1623 | name = "serde_derive_internals" 1624 | version = "0.26.0" 1625 | source = "registry+https://github.com/rust-lang/crates.io-index" 1626 | checksum = "85bf8229e7920a9f636479437026331ce11aa132b4dde37d121944a44d6e5f3c" 1627 | dependencies = [ 1628 | "proc-macro2", 1629 | "quote", 1630 | "syn 1.0.109", 1631 | ] 1632 | 1633 | [[package]] 1634 | name = "serde_json" 1635 | version = "1.0.96" 1636 | source = "registry+https://github.com/rust-lang/crates.io-index" 1637 | checksum = "057d394a50403bcac12672b2b18fb387ab6d289d957dab67dd201875391e52f1" 1638 | dependencies = [ 1639 | "itoa", 1640 | "ryu", 1641 | "serde", 1642 | ] 1643 | 1644 | [[package]] 1645 | name = "serde_path_to_error" 1646 | version = "0.1.11" 1647 | source = "registry+https://github.com/rust-lang/crates.io-index" 1648 | checksum = "f7f05c1d5476066defcdfacce1f52fc3cae3af1d3089727100c02ae92e5abbe0" 1649 | dependencies = [ 1650 | "serde", 1651 | ] 1652 | 1653 | [[package]] 1654 | name = "serde_qs" 1655 | version = "0.12.0" 1656 | source = "registry+https://github.com/rust-lang/crates.io-index" 1657 | checksum = "0431a35568651e363364210c91983c1da5eb29404d9f0928b67d4ebcfa7d330c" 1658 | dependencies = [ 1659 | "axum", 1660 | "futures", 1661 | "percent-encoding", 1662 | "serde", 1663 | "thiserror", 1664 | ] 1665 | 1666 | [[package]] 1667 | name = "serde_urlencoded" 1668 | version = "0.7.1" 1669 | source = "registry+https://github.com/rust-lang/crates.io-index" 1670 | checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" 1671 | dependencies = [ 1672 | "form_urlencoded", 1673 | "itoa", 1674 | "ryu", 1675 | "serde", 1676 | ] 1677 | 1678 | [[package]] 1679 | name = "sha1" 1680 | version = "0.10.5" 1681 | source = "registry+https://github.com/rust-lang/crates.io-index" 1682 | checksum = "f04293dc80c3993519f2d7f6f511707ee7094fe0c6d3406feb330cdb3540eba3" 1683 | dependencies = [ 1684 | "cfg-if", 1685 | "cpufeatures", 1686 | "digest", 1687 | ] 1688 | 1689 | [[package]] 1690 | name = "sharded-slab" 1691 | version = "0.1.4" 1692 | source = "registry+https://github.com/rust-lang/crates.io-index" 1693 | checksum = "900fba806f70c630b0a382d0d825e17a0f19fcd059a2ade1ff237bcddf446b31" 1694 | dependencies = [ 1695 | "lazy_static", 1696 | ] 1697 | 1698 | [[package]] 1699 | name = "signal-hook-registry" 1700 | version = "1.4.1" 1701 | source = "registry+https://github.com/rust-lang/crates.io-index" 1702 | checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" 1703 | dependencies = [ 1704 | "libc", 1705 | ] 1706 | 1707 | [[package]] 1708 | name = "slab" 1709 | version = "0.4.8" 1710 | source = "registry+https://github.com/rust-lang/crates.io-index" 1711 | checksum = "6528351c9bc8ab22353f9d776db39a20288e8d6c37ef8cfe3317cf875eecfc2d" 1712 | dependencies = [ 1713 | "autocfg", 1714 | ] 1715 | 1716 | [[package]] 1717 | name = "smallvec" 1718 | version = "1.10.0" 1719 | source = "registry+https://github.com/rust-lang/crates.io-index" 1720 | checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" 1721 | 1722 | [[package]] 1723 | name = "socket2" 1724 | version = "0.4.9" 1725 | source = "registry+https://github.com/rust-lang/crates.io-index" 1726 | checksum = "64a4a911eed85daf18834cfaa86a79b7d266ff93ff5ba14005426219480ed662" 1727 | dependencies = [ 1728 | "libc", 1729 | "winapi", 1730 | ] 1731 | 1732 | [[package]] 1733 | name = "socket2" 1734 | version = "0.5.5" 1735 | source = "registry+https://github.com/rust-lang/crates.io-index" 1736 | checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" 1737 | dependencies = [ 1738 | "libc", 1739 | "windows-sys 0.48.0", 1740 | ] 1741 | 1742 | [[package]] 1743 | name = "spin" 1744 | version = "0.9.8" 1745 | source = "registry+https://github.com/rust-lang/crates.io-index" 1746 | checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" 1747 | dependencies = [ 1748 | "lock_api", 1749 | ] 1750 | 1751 | [[package]] 1752 | name = "strsim" 1753 | version = "0.10.0" 1754 | source = "registry+https://github.com/rust-lang/crates.io-index" 1755 | checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" 1756 | 1757 | [[package]] 1758 | name = "syn" 1759 | version = "1.0.109" 1760 | source = "registry+https://github.com/rust-lang/crates.io-index" 1761 | checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" 1762 | dependencies = [ 1763 | "proc-macro2", 1764 | "quote", 1765 | "unicode-ident", 1766 | ] 1767 | 1768 | [[package]] 1769 | name = "syn" 1770 | version = "2.0.18" 1771 | source = "registry+https://github.com/rust-lang/crates.io-index" 1772 | checksum = "32d41677bcbe24c20c52e7c70b0d8db04134c5d1066bf98662e2871ad200ea3e" 1773 | dependencies = [ 1774 | "proc-macro2", 1775 | "quote", 1776 | "unicode-ident", 1777 | ] 1778 | 1779 | [[package]] 1780 | name = "sync_wrapper" 1781 | version = "0.1.2" 1782 | source = "registry+https://github.com/rust-lang/crates.io-index" 1783 | checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" 1784 | 1785 | [[package]] 1786 | name = "tempfile" 1787 | version = "3.6.0" 1788 | source = "registry+https://github.com/rust-lang/crates.io-index" 1789 | checksum = "31c0432476357e58790aaa47a8efb0c5138f137343f3b5f23bd36a27e3b0a6d6" 1790 | dependencies = [ 1791 | "autocfg", 1792 | "cfg-if", 1793 | "fastrand", 1794 | "redox_syscall", 1795 | "rustix", 1796 | "windows-sys 0.48.0", 1797 | ] 1798 | 1799 | [[package]] 1800 | name = "thiserror" 1801 | version = "1.0.40" 1802 | source = "registry+https://github.com/rust-lang/crates.io-index" 1803 | checksum = "978c9a314bd8dc99be594bc3c175faaa9794be04a5a5e153caba6915336cebac" 1804 | dependencies = [ 1805 | "thiserror-impl", 1806 | ] 1807 | 1808 | [[package]] 1809 | name = "thiserror-impl" 1810 | version = "1.0.40" 1811 | source = "registry+https://github.com/rust-lang/crates.io-index" 1812 | checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f" 1813 | dependencies = [ 1814 | "proc-macro2", 1815 | "quote", 1816 | "syn 2.0.18", 1817 | ] 1818 | 1819 | [[package]] 1820 | name = "thread_local" 1821 | version = "1.1.7" 1822 | source = "registry+https://github.com/rust-lang/crates.io-index" 1823 | checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152" 1824 | dependencies = [ 1825 | "cfg-if", 1826 | "once_cell", 1827 | ] 1828 | 1829 | [[package]] 1830 | name = "time" 1831 | version = "0.1.45" 1832 | source = "registry+https://github.com/rust-lang/crates.io-index" 1833 | checksum = "1b797afad3f312d1c66a56d11d0316f916356d11bd158fbc6ca6389ff6bf805a" 1834 | dependencies = [ 1835 | "libc", 1836 | "wasi 0.10.0+wasi-snapshot-preview1", 1837 | "winapi", 1838 | ] 1839 | 1840 | [[package]] 1841 | name = "time" 1842 | version = "0.3.22" 1843 | source = "registry+https://github.com/rust-lang/crates.io-index" 1844 | checksum = "ea9e1b3cf1243ae005d9e74085d4d542f3125458f3a81af210d901dcd7411efd" 1845 | dependencies = [ 1846 | "serde", 1847 | "time-core", 1848 | "time-macros", 1849 | ] 1850 | 1851 | [[package]] 1852 | name = "time-core" 1853 | version = "0.1.1" 1854 | source = "registry+https://github.com/rust-lang/crates.io-index" 1855 | checksum = "7300fbefb4dadc1af235a9cef3737cea692a9d97e1b9cbcd4ebdae6f8868e6fb" 1856 | 1857 | [[package]] 1858 | name = "time-macros" 1859 | version = "0.2.9" 1860 | source = "registry+https://github.com/rust-lang/crates.io-index" 1861 | checksum = "372950940a5f07bf38dbe211d7283c9e6d7327df53794992d293e534c733d09b" 1862 | dependencies = [ 1863 | "time-core", 1864 | ] 1865 | 1866 | [[package]] 1867 | name = "tinyvec" 1868 | version = "1.6.0" 1869 | source = "registry+https://github.com/rust-lang/crates.io-index" 1870 | checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50" 1871 | dependencies = [ 1872 | "tinyvec_macros", 1873 | ] 1874 | 1875 | [[package]] 1876 | name = "tinyvec_macros" 1877 | version = "0.1.1" 1878 | source = "registry+https://github.com/rust-lang/crates.io-index" 1879 | checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" 1880 | 1881 | [[package]] 1882 | name = "titlecase" 1883 | version = "2.2.1" 1884 | source = "registry+https://github.com/rust-lang/crates.io-index" 1885 | checksum = "38397a8cdb017cfeb48bf6c154d6de975ac69ffeed35980fde199d2ee0842042" 1886 | dependencies = [ 1887 | "joinery", 1888 | "lazy_static", 1889 | "regex", 1890 | ] 1891 | 1892 | [[package]] 1893 | name = "tokio" 1894 | version = "1.35.0" 1895 | source = "registry+https://github.com/rust-lang/crates.io-index" 1896 | checksum = "841d45b238a16291a4e1584e61820b8ae57d696cc5015c459c229ccc6990cc1c" 1897 | dependencies = [ 1898 | "backtrace", 1899 | "bytes", 1900 | "libc", 1901 | "mio", 1902 | "num_cpus", 1903 | "parking_lot", 1904 | "pin-project-lite", 1905 | "signal-hook-registry", 1906 | "socket2 0.5.5", 1907 | "tokio-macros", 1908 | "windows-sys 0.48.0", 1909 | ] 1910 | 1911 | [[package]] 1912 | name = "tokio-macros" 1913 | version = "2.2.0" 1914 | source = "registry+https://github.com/rust-lang/crates.io-index" 1915 | checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" 1916 | dependencies = [ 1917 | "proc-macro2", 1918 | "quote", 1919 | "syn 2.0.18", 1920 | ] 1921 | 1922 | [[package]] 1923 | name = "tokio-native-tls" 1924 | version = "0.3.1" 1925 | source = "registry+https://github.com/rust-lang/crates.io-index" 1926 | checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" 1927 | dependencies = [ 1928 | "native-tls", 1929 | "tokio", 1930 | ] 1931 | 1932 | [[package]] 1933 | name = "tokio-util" 1934 | version = "0.7.8" 1935 | source = "registry+https://github.com/rust-lang/crates.io-index" 1936 | checksum = "806fe8c2c87eccc8b3267cbae29ed3ab2d0bd37fca70ab622e46aaa9375ddb7d" 1937 | dependencies = [ 1938 | "bytes", 1939 | "futures-core", 1940 | "futures-sink", 1941 | "pin-project-lite", 1942 | "tokio", 1943 | "tracing", 1944 | ] 1945 | 1946 | [[package]] 1947 | name = "tower" 1948 | version = "0.4.13" 1949 | source = "registry+https://github.com/rust-lang/crates.io-index" 1950 | checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" 1951 | dependencies = [ 1952 | "futures-core", 1953 | "futures-util", 1954 | "pin-project", 1955 | "pin-project-lite", 1956 | "tokio", 1957 | "tower-layer", 1958 | "tower-service", 1959 | "tracing", 1960 | ] 1961 | 1962 | [[package]] 1963 | name = "tower-layer" 1964 | version = "0.3.2" 1965 | source = "registry+https://github.com/rust-lang/crates.io-index" 1966 | checksum = "c20c8dbed6283a09604c3e69b4b7eeb54e298b8a600d4d5ecb5ad39de609f1d0" 1967 | 1968 | [[package]] 1969 | name = "tower-service" 1970 | version = "0.3.2" 1971 | source = "registry+https://github.com/rust-lang/crates.io-index" 1972 | checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52" 1973 | 1974 | [[package]] 1975 | name = "tracing" 1976 | version = "0.1.37" 1977 | source = "registry+https://github.com/rust-lang/crates.io-index" 1978 | checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8" 1979 | dependencies = [ 1980 | "cfg-if", 1981 | "log", 1982 | "pin-project-lite", 1983 | "tracing-attributes", 1984 | "tracing-core", 1985 | ] 1986 | 1987 | [[package]] 1988 | name = "tracing-attributes" 1989 | version = "0.1.24" 1990 | source = "registry+https://github.com/rust-lang/crates.io-index" 1991 | checksum = "0f57e3ca2a01450b1a921183a9c9cbfda207fd822cef4ccb00a65402cbba7a74" 1992 | dependencies = [ 1993 | "proc-macro2", 1994 | "quote", 1995 | "syn 2.0.18", 1996 | ] 1997 | 1998 | [[package]] 1999 | name = "tracing-core" 2000 | version = "0.1.31" 2001 | source = "registry+https://github.com/rust-lang/crates.io-index" 2002 | checksum = "0955b8137a1df6f1a2e9a37d8a6656291ff0297c1a97c24e0d8425fe2312f79a" 2003 | dependencies = [ 2004 | "once_cell", 2005 | "valuable", 2006 | ] 2007 | 2008 | [[package]] 2009 | name = "tracing-log" 2010 | version = "0.1.3" 2011 | source = "registry+https://github.com/rust-lang/crates.io-index" 2012 | checksum = "78ddad33d2d10b1ed7eb9d1f518a5674713876e97e5bb9b7345a7984fbb4f922" 2013 | dependencies = [ 2014 | "lazy_static", 2015 | "log", 2016 | "tracing-core", 2017 | ] 2018 | 2019 | [[package]] 2020 | name = "tracing-subscriber" 2021 | version = "0.3.17" 2022 | source = "registry+https://github.com/rust-lang/crates.io-index" 2023 | checksum = "30a651bc37f915e81f087d86e62a18eec5f79550c7faff886f7090b4ea757c77" 2024 | dependencies = [ 2025 | "matchers", 2026 | "nu-ansi-term", 2027 | "once_cell", 2028 | "regex", 2029 | "sharded-slab", 2030 | "smallvec", 2031 | "thread_local", 2032 | "tracing", 2033 | "tracing-core", 2034 | "tracing-log", 2035 | ] 2036 | 2037 | [[package]] 2038 | name = "tree_magic_db" 2039 | version = "3.0.0" 2040 | source = "registry+https://github.com/rust-lang/crates.io-index" 2041 | checksum = "e73fc24a5427b3b15e2b0bcad8ef61b5affb1da8ac89c8bf3f196c8692d57f02" 2042 | 2043 | [[package]] 2044 | name = "tree_magic_mini" 2045 | version = "3.0.3" 2046 | source = "registry+https://github.com/rust-lang/crates.io-index" 2047 | checksum = "91adfd0607cacf6e4babdb870e9bec4037c1c4b151cfd279ccefc5e0c7feaa6d" 2048 | dependencies = [ 2049 | "bytecount", 2050 | "fnv", 2051 | "lazy_static", 2052 | "nom", 2053 | "once_cell", 2054 | "petgraph", 2055 | "tree_magic_db", 2056 | ] 2057 | 2058 | [[package]] 2059 | name = "try-lock" 2060 | version = "0.2.4" 2061 | source = "registry+https://github.com/rust-lang/crates.io-index" 2062 | checksum = "3528ecfd12c466c6f163363caf2d02a71161dd5e1cc6ae7b34207ea2d42d81ed" 2063 | 2064 | [[package]] 2065 | name = "typenum" 2066 | version = "1.16.0" 2067 | source = "registry+https://github.com/rust-lang/crates.io-index" 2068 | checksum = "497961ef93d974e23eb6f433eb5fe1b7930b659f06d12dec6fc44a8f554c0bba" 2069 | 2070 | [[package]] 2071 | name = "unicase" 2072 | version = "2.6.0" 2073 | source = "registry+https://github.com/rust-lang/crates.io-index" 2074 | checksum = "50f37be617794602aabbeee0be4f259dc1778fabe05e2d67ee8f79326d5cb4f6" 2075 | dependencies = [ 2076 | "version_check", 2077 | ] 2078 | 2079 | [[package]] 2080 | name = "unicode-bidi" 2081 | version = "0.3.13" 2082 | source = "registry+https://github.com/rust-lang/crates.io-index" 2083 | checksum = "92888ba5573ff080736b3648696b70cafad7d250551175acbaa4e0385b3e1460" 2084 | 2085 | [[package]] 2086 | name = "unicode-ident" 2087 | version = "1.0.9" 2088 | source = "registry+https://github.com/rust-lang/crates.io-index" 2089 | checksum = "b15811caf2415fb889178633e7724bad2509101cde276048e013b9def5e51fa0" 2090 | 2091 | [[package]] 2092 | name = "unicode-normalization" 2093 | version = "0.1.22" 2094 | source = "registry+https://github.com/rust-lang/crates.io-index" 2095 | checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921" 2096 | dependencies = [ 2097 | "tinyvec", 2098 | ] 2099 | 2100 | [[package]] 2101 | name = "url" 2102 | version = "2.4.0" 2103 | source = "registry+https://github.com/rust-lang/crates.io-index" 2104 | checksum = "50bff7831e19200a85b17131d085c25d7811bc4e186efdaf54bbd132994a88cb" 2105 | dependencies = [ 2106 | "form_urlencoded", 2107 | "idna", 2108 | "percent-encoding", 2109 | "serde", 2110 | ] 2111 | 2112 | [[package]] 2113 | name = "utf8parse" 2114 | version = "0.2.1" 2115 | source = "registry+https://github.com/rust-lang/crates.io-index" 2116 | checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" 2117 | 2118 | [[package]] 2119 | name = "uuid" 2120 | version = "1.3.3" 2121 | source = "registry+https://github.com/rust-lang/crates.io-index" 2122 | checksum = "345444e32442451b267fc254ae85a209c64be56d2890e601a0c37ff0c3c5ecd2" 2123 | dependencies = [ 2124 | "getrandom", 2125 | ] 2126 | 2127 | [[package]] 2128 | name = "valuable" 2129 | version = "0.1.0" 2130 | source = "registry+https://github.com/rust-lang/crates.io-index" 2131 | checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" 2132 | 2133 | [[package]] 2134 | name = "vcpkg" 2135 | version = "0.2.15" 2136 | source = "registry+https://github.com/rust-lang/crates.io-index" 2137 | checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" 2138 | 2139 | [[package]] 2140 | name = "version_check" 2141 | version = "0.9.4" 2142 | source = "registry+https://github.com/rust-lang/crates.io-index" 2143 | checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" 2144 | 2145 | [[package]] 2146 | name = "want" 2147 | version = "0.3.0" 2148 | source = "registry+https://github.com/rust-lang/crates.io-index" 2149 | checksum = "1ce8a968cb1cd110d136ff8b819a556d6fb6d919363c61534f6860c7eb172ba0" 2150 | dependencies = [ 2151 | "log", 2152 | "try-lock", 2153 | ] 2154 | 2155 | [[package]] 2156 | name = "wasi" 2157 | version = "0.10.0+wasi-snapshot-preview1" 2158 | source = "registry+https://github.com/rust-lang/crates.io-index" 2159 | checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f" 2160 | 2161 | [[package]] 2162 | name = "wasi" 2163 | version = "0.11.0+wasi-snapshot-preview1" 2164 | source = "registry+https://github.com/rust-lang/crates.io-index" 2165 | checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" 2166 | 2167 | [[package]] 2168 | name = "wasm-bindgen" 2169 | version = "0.2.86" 2170 | source = "registry+https://github.com/rust-lang/crates.io-index" 2171 | checksum = "5bba0e8cb82ba49ff4e229459ff22a191bbe9a1cb3a341610c9c33efc27ddf73" 2172 | dependencies = [ 2173 | "cfg-if", 2174 | "wasm-bindgen-macro", 2175 | ] 2176 | 2177 | [[package]] 2178 | name = "wasm-bindgen-backend" 2179 | version = "0.2.86" 2180 | source = "registry+https://github.com/rust-lang/crates.io-index" 2181 | checksum = "19b04bc93f9d6bdee709f6bd2118f57dd6679cf1176a1af464fca3ab0d66d8fb" 2182 | dependencies = [ 2183 | "bumpalo", 2184 | "log", 2185 | "once_cell", 2186 | "proc-macro2", 2187 | "quote", 2188 | "syn 2.0.18", 2189 | "wasm-bindgen-shared", 2190 | ] 2191 | 2192 | [[package]] 2193 | name = "wasm-bindgen-futures" 2194 | version = "0.4.36" 2195 | source = "registry+https://github.com/rust-lang/crates.io-index" 2196 | checksum = "2d1985d03709c53167ce907ff394f5316aa22cb4e12761295c5dc57dacb6297e" 2197 | dependencies = [ 2198 | "cfg-if", 2199 | "js-sys", 2200 | "wasm-bindgen", 2201 | "web-sys", 2202 | ] 2203 | 2204 | [[package]] 2205 | name = "wasm-bindgen-macro" 2206 | version = "0.2.86" 2207 | source = "registry+https://github.com/rust-lang/crates.io-index" 2208 | checksum = "14d6b024f1a526bb0234f52840389927257beb670610081360e5a03c5df9c258" 2209 | dependencies = [ 2210 | "quote", 2211 | "wasm-bindgen-macro-support", 2212 | ] 2213 | 2214 | [[package]] 2215 | name = "wasm-bindgen-macro-support" 2216 | version = "0.2.86" 2217 | source = "registry+https://github.com/rust-lang/crates.io-index" 2218 | checksum = "e128beba882dd1eb6200e1dc92ae6c5dbaa4311aa7bb211ca035779e5efc39f8" 2219 | dependencies = [ 2220 | "proc-macro2", 2221 | "quote", 2222 | "syn 2.0.18", 2223 | "wasm-bindgen-backend", 2224 | "wasm-bindgen-shared", 2225 | ] 2226 | 2227 | [[package]] 2228 | name = "wasm-bindgen-shared" 2229 | version = "0.2.86" 2230 | source = "registry+https://github.com/rust-lang/crates.io-index" 2231 | checksum = "ed9d5b4305409d1fc9482fee2d7f9bcbf24b3972bf59817ef757e23982242a93" 2232 | 2233 | [[package]] 2234 | name = "web-sys" 2235 | version = "0.3.63" 2236 | source = "registry+https://github.com/rust-lang/crates.io-index" 2237 | checksum = "3bdd9ef4e984da1187bf8110c5cf5b845fbc87a23602cdf912386a76fcd3a7c2" 2238 | dependencies = [ 2239 | "js-sys", 2240 | "wasm-bindgen", 2241 | ] 2242 | 2243 | [[package]] 2244 | name = "winapi" 2245 | version = "0.3.9" 2246 | source = "registry+https://github.com/rust-lang/crates.io-index" 2247 | checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" 2248 | dependencies = [ 2249 | "winapi-i686-pc-windows-gnu", 2250 | "winapi-x86_64-pc-windows-gnu", 2251 | ] 2252 | 2253 | [[package]] 2254 | name = "winapi-i686-pc-windows-gnu" 2255 | version = "0.4.0" 2256 | source = "registry+https://github.com/rust-lang/crates.io-index" 2257 | checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" 2258 | 2259 | [[package]] 2260 | name = "winapi-x86_64-pc-windows-gnu" 2261 | version = "0.4.0" 2262 | source = "registry+https://github.com/rust-lang/crates.io-index" 2263 | checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" 2264 | 2265 | [[package]] 2266 | name = "windows" 2267 | version = "0.48.0" 2268 | source = "registry+https://github.com/rust-lang/crates.io-index" 2269 | checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" 2270 | dependencies = [ 2271 | "windows-targets", 2272 | ] 2273 | 2274 | [[package]] 2275 | name = "windows-sys" 2276 | version = "0.42.0" 2277 | source = "registry+https://github.com/rust-lang/crates.io-index" 2278 | checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" 2279 | dependencies = [ 2280 | "windows_aarch64_gnullvm 0.42.2", 2281 | "windows_aarch64_msvc 0.42.2", 2282 | "windows_i686_gnu 0.42.2", 2283 | "windows_i686_msvc 0.42.2", 2284 | "windows_x86_64_gnu 0.42.2", 2285 | "windows_x86_64_gnullvm 0.42.2", 2286 | "windows_x86_64_msvc 0.42.2", 2287 | ] 2288 | 2289 | [[package]] 2290 | name = "windows-sys" 2291 | version = "0.48.0" 2292 | source = "registry+https://github.com/rust-lang/crates.io-index" 2293 | checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" 2294 | dependencies = [ 2295 | "windows-targets", 2296 | ] 2297 | 2298 | [[package]] 2299 | name = "windows-targets" 2300 | version = "0.48.0" 2301 | source = "registry+https://github.com/rust-lang/crates.io-index" 2302 | checksum = "7b1eb6f0cd7c80c79759c929114ef071b87354ce476d9d94271031c0497adfd5" 2303 | dependencies = [ 2304 | "windows_aarch64_gnullvm 0.48.0", 2305 | "windows_aarch64_msvc 0.48.0", 2306 | "windows_i686_gnu 0.48.0", 2307 | "windows_i686_msvc 0.48.0", 2308 | "windows_x86_64_gnu 0.48.0", 2309 | "windows_x86_64_gnullvm 0.48.0", 2310 | "windows_x86_64_msvc 0.48.0", 2311 | ] 2312 | 2313 | [[package]] 2314 | name = "windows_aarch64_gnullvm" 2315 | version = "0.42.2" 2316 | source = "registry+https://github.com/rust-lang/crates.io-index" 2317 | checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" 2318 | 2319 | [[package]] 2320 | name = "windows_aarch64_gnullvm" 2321 | version = "0.48.0" 2322 | source = "registry+https://github.com/rust-lang/crates.io-index" 2323 | checksum = "91ae572e1b79dba883e0d315474df7305d12f569b400fcf90581b06062f7e1bc" 2324 | 2325 | [[package]] 2326 | name = "windows_aarch64_msvc" 2327 | version = "0.42.2" 2328 | source = "registry+https://github.com/rust-lang/crates.io-index" 2329 | checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" 2330 | 2331 | [[package]] 2332 | name = "windows_aarch64_msvc" 2333 | version = "0.48.0" 2334 | source = "registry+https://github.com/rust-lang/crates.io-index" 2335 | checksum = "b2ef27e0d7bdfcfc7b868b317c1d32c641a6fe4629c171b8928c7b08d98d7cf3" 2336 | 2337 | [[package]] 2338 | name = "windows_i686_gnu" 2339 | version = "0.42.2" 2340 | source = "registry+https://github.com/rust-lang/crates.io-index" 2341 | checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" 2342 | 2343 | [[package]] 2344 | name = "windows_i686_gnu" 2345 | version = "0.48.0" 2346 | source = "registry+https://github.com/rust-lang/crates.io-index" 2347 | checksum = "622a1962a7db830d6fd0a69683c80a18fda201879f0f447f065a3b7467daa241" 2348 | 2349 | [[package]] 2350 | name = "windows_i686_msvc" 2351 | version = "0.42.2" 2352 | source = "registry+https://github.com/rust-lang/crates.io-index" 2353 | checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" 2354 | 2355 | [[package]] 2356 | name = "windows_i686_msvc" 2357 | version = "0.48.0" 2358 | source = "registry+https://github.com/rust-lang/crates.io-index" 2359 | checksum = "4542c6e364ce21bf45d69fdd2a8e455fa38d316158cfd43b3ac1c5b1b19f8e00" 2360 | 2361 | [[package]] 2362 | name = "windows_x86_64_gnu" 2363 | version = "0.42.2" 2364 | source = "registry+https://github.com/rust-lang/crates.io-index" 2365 | checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" 2366 | 2367 | [[package]] 2368 | name = "windows_x86_64_gnu" 2369 | version = "0.48.0" 2370 | source = "registry+https://github.com/rust-lang/crates.io-index" 2371 | checksum = "ca2b8a661f7628cbd23440e50b05d705db3686f894fc9580820623656af974b1" 2372 | 2373 | [[package]] 2374 | name = "windows_x86_64_gnullvm" 2375 | version = "0.42.2" 2376 | source = "registry+https://github.com/rust-lang/crates.io-index" 2377 | checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" 2378 | 2379 | [[package]] 2380 | name = "windows_x86_64_gnullvm" 2381 | version = "0.48.0" 2382 | source = "registry+https://github.com/rust-lang/crates.io-index" 2383 | checksum = "7896dbc1f41e08872e9d5e8f8baa8fdd2677f29468c4e156210174edc7f7b953" 2384 | 2385 | [[package]] 2386 | name = "windows_x86_64_msvc" 2387 | version = "0.42.2" 2388 | source = "registry+https://github.com/rust-lang/crates.io-index" 2389 | checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" 2390 | 2391 | [[package]] 2392 | name = "windows_x86_64_msvc" 2393 | version = "0.48.0" 2394 | source = "registry+https://github.com/rust-lang/crates.io-index" 2395 | checksum = "1a515f5799fe4961cb532f983ce2b23082366b898e52ffbce459c86f67c8378a" 2396 | 2397 | [[package]] 2398 | name = "winreg" 2399 | version = "0.10.1" 2400 | source = "registry+https://github.com/rust-lang/crates.io-index" 2401 | checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" 2402 | dependencies = [ 2403 | "winapi", 2404 | ] 2405 | -------------------------------------------------------------------------------- /examples/hello-world/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-hello-world" 3 | version = "0.0.0" 4 | edition = "2021" 5 | publish = false 6 | authors = ["Miguel Piedrafita "] 7 | 8 | [package.metadata.cog] 9 | image = "hello-world-rs" 10 | 11 | [dependencies] 12 | serde = "1.0.163" 13 | anyhow = "1.0.71" 14 | schemars = "0.8.12" 15 | cog-rust = { path = "../../lib" } 16 | tokio = { version = "1.28.2", features = ["full"] } 17 | -------------------------------------------------------------------------------- /examples/hello-world/src/main.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use cog_rust::Cog; 3 | use schemars::JsonSchema; 4 | 5 | #[derive(serde::Deserialize, JsonSchema)] 6 | struct ModelRequest { 7 | /// Text to prefix with 'hello ' 8 | text: String, 9 | } 10 | 11 | struct ExampleModel { 12 | prefix: String, 13 | } 14 | 15 | impl Cog for ExampleModel { 16 | type Request = ModelRequest; 17 | type Response = String; 18 | 19 | async fn setup() -> Result { 20 | Ok(Self { 21 | prefix: "hello".to_string(), 22 | }) 23 | } 24 | 25 | fn predict(&self, input: Self::Request) -> Result { 26 | Ok(format!("{} {}", self.prefix, input.text)) 27 | } 28 | } 29 | 30 | cog_rust::start!(ExampleModel); 31 | -------------------------------------------------------------------------------- /examples/resnet/.gitignore: -------------------------------------------------------------------------------- 1 | weights/model.safetensors 2 | -------------------------------------------------------------------------------- /examples/resnet/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-resnet" 3 | version = "0.0.0" 4 | edition = "2021" 5 | publish = false 6 | authors = ["Miguel Piedrafita "] 7 | 8 | [package.metadata.cog] 9 | image = "resnet-rs" 10 | 11 | [dependencies] 12 | serde = "1.0.163" 13 | anyhow = "1.0.71" 14 | schemars = "0.8.12" 15 | cog-rust = { path = "../../lib" } 16 | tokio = { version = "1.28.2", features = ["full"] } 17 | tch = { version = "0.13.0", features = ["download-libtorch"] } 18 | -------------------------------------------------------------------------------- /examples/resnet/README.md: -------------------------------------------------------------------------------- 1 | # resnet 2 | 3 | This model classifies images. 4 | 5 | --- 6 | 7 | Download the pre-trained weights: 8 | 9 | ``` 10 | curl -L -o weights/model.safetensors https://huggingface.co/microsoft/resnet-50/resolve/refs%2Fpr%2F4/model.safetensors 11 | ``` 12 | 13 | Build the image: 14 | 15 | ```sh 16 | cargo cog build 17 | ``` 18 | 19 | Now you can run predictions on the model: 20 | 21 | ```sh 22 | cargo cog predict -i image=@cat.png 23 | ``` 24 | -------------------------------------------------------------------------------- /examples/resnet/src/main.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use cog_rust::Cog; 3 | use schemars::JsonSchema; 4 | use std::collections::HashMap; 5 | use tch::{ 6 | nn::{ModuleT, VarStore}, 7 | vision::{imagenet, resnet::resnet50}, 8 | Device, 9 | }; 10 | 11 | #[derive(serde::Deserialize, JsonSchema)] 12 | struct ModelRequest { 13 | /// Image to classify 14 | image: cog_rust::Path, 15 | } 16 | 17 | struct ResnetModel { 18 | model: Box, 19 | } 20 | 21 | impl Cog for ResnetModel { 22 | type Request = ModelRequest; 23 | type Response = HashMap; 24 | 25 | async fn setup() -> Result { 26 | let mut vs = VarStore::new(Device::Cpu); 27 | vs.load("weights/model.safetensors")?; 28 | let model = Box::new(resnet50(&vs.root(), imagenet::CLASS_COUNT)); 29 | 30 | Ok(Self { model }) 31 | } 32 | 33 | fn predict(&self, input: Self::Request) -> Result { 34 | let image = imagenet::load_image_and_resize224(&input.image)?; 35 | let output = self 36 | .model 37 | .forward_t(&image.unsqueeze(0), false) 38 | .softmax(-1, tch::Kind::Float); 39 | 40 | Ok(imagenet::top(&output, 5) 41 | .into_iter() 42 | .map(|(prob, class)| (class, 100.0 * prob)) 43 | .collect()) 44 | } 45 | } 46 | 47 | cog_rust::start!(ResnetModel); 48 | -------------------------------------------------------------------------------- /examples/resnet/weights/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/m1guelpf/cog-rust/4d6666a2245a6ae9a6e14973f3a430a94e3fd4f9/examples/resnet/weights/.gitkeep -------------------------------------------------------------------------------- /examples/stable-diffusion/.gitignore: -------------------------------------------------------------------------------- 1 | weights/* 2 | !weights/.gitkeep 3 | -------------------------------------------------------------------------------- /examples/stable-diffusion/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-stable-diffusion" 3 | version = "0.0.0" 4 | edition = "2021" 5 | publish = false 6 | authors = ["Miguel Piedrafita "] 7 | 8 | [package.metadata.cog] 9 | gpu = true 10 | image = "stable-diffusion-rs" 11 | 12 | [dependencies] 13 | serde = "1.0.163" 14 | anyhow = "1.0.71" 15 | schemars = "0.8.12" 16 | diffusers = "0.3.1" 17 | async-trait = "0.1.68" 18 | cog-rust = { path = "../../lib" } 19 | tokio = { version = "1.28.2", features = ["full"] } 20 | tch = { version = "0.13.0", features = ["download-libtorch"] } 21 | -------------------------------------------------------------------------------- /examples/stable-diffusion/README.md: -------------------------------------------------------------------------------- 1 | # Stable Diffusion 2 | 3 | This is an implementation of the Diffusers Stable Diffusion v2.1 as a Cog model. 4 | 5 | --- 6 | 7 | Download the pre-trained weights: 8 | 9 | ``` 10 | curl -L -o weights/bpe_simple_vocab_16e6.txt https://huggingface.co/lmz/rust-stable-diffusion-v2-1/raw/main/weights/bpe_simple_vocab_16e6.txt 11 | curl -L -o weights/clip_v2.1.safetensors https://huggingface.co/lmz/rust-stable-diffusion-v2-1/resolve/main/weights/clip_v2.1.safetensors 12 | curl -L -o weights/unet_v2.1.safetensors https://huggingface.co/lmz/rust-stable-diffusion-v2-1/resolve/main/weights/unet_v2.1.safetensors 13 | curl -L -o weights/vae_v2.1.safetensors https://huggingface.co/lmz/rust-stable-diffusion-v2-1/resolve/main/weights/vae_v2.1.safetensors 14 | ``` 15 | 16 | Build the image: 17 | 18 | ```sh 19 | cargo cog build 20 | ``` 21 | 22 | Now you can run predictions on the model: 23 | 24 | ```sh 25 | cargo cog predict -i "prompt=a photo of a cat" 26 | ``` 27 | -------------------------------------------------------------------------------- /examples/stable-diffusion/src/main.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use cog_rust::{Cog, Path}; 3 | use diffusers::{ 4 | models::{unet_2d::UNet2DConditionModel, vae::AutoEncoderKL}, 5 | pipelines::stable_diffusion::{self, StableDiffusionConfig}, 6 | transformers::clip::{self, Tokenizer}, 7 | utils::DeviceSetup, 8 | }; 9 | use schemars::JsonSchema; 10 | use std::path::PathBuf; 11 | use tch::{nn::Module, Device, Kind, Tensor}; 12 | 13 | #[derive(serde::Deserialize, JsonSchema)] 14 | struct ModelRequest { 15 | /// Input prompt 16 | prompt: String, 17 | 18 | /// Random seed. Leave blank to randomize the seed 19 | seed: Option, 20 | 21 | /// Number of images to output. 22 | #[validate(range(min = 1, max = 4))] 23 | num_outputs: Option, 24 | 25 | /// Number of denoising steps. 26 | #[validate(range(min = 1, max = 500))] 27 | num_inference_steps: Option, 28 | 29 | /// Scale for classifier-free guidance. 30 | #[validate(range(min = 1, max = 20))] 31 | guidance_scale: Option, 32 | } 33 | 34 | struct StableDiffusion { 35 | vae: AutoEncoderKL, 36 | tokenizer: Tokenizer, 37 | devices: DeviceSetup, 38 | unet: UNet2DConditionModel, 39 | sd_config: StableDiffusionConfig, 40 | text_model: clip::ClipTextTransformer, 41 | } 42 | 43 | impl Cog for StableDiffusion { 44 | type Request = ModelRequest; 45 | type Response = Vec; 46 | 47 | async fn setup() -> Result { 48 | tch::maybe_init_cuda(); 49 | let sd_config = stable_diffusion::StableDiffusionConfig::v2_1(None, None, None); 50 | let device_setup = diffusers::utils::DeviceSetup::new(vec![]); 51 | let tokenizer = 52 | clip::Tokenizer::create("weights/bpe_simple_vocab_16e6.txt", &sd_config.clip)?; 53 | let text_model = sd_config 54 | .build_clip_transformer("weights/clip_v2.1.safetensors", device_setup.get("clip"))?; 55 | let vae = sd_config.build_vae("weights/vae_v2.1.safetensors", device_setup.get("vae"))?; 56 | let unet = 57 | sd_config.build_unet("weights/unet_v2.1.safetensors", device_setup.get("unet"), 4)?; 58 | 59 | Ok(Self { 60 | vae, 61 | unet, 62 | tokenizer, 63 | sd_config, 64 | text_model, 65 | devices: device_setup, 66 | }) 67 | } 68 | 69 | fn predict(&self, input: Self::Request) -> Result { 70 | let _no_grad_guard = tch::no_grad_guard(); 71 | let scheduler = self 72 | .sd_config 73 | .build_scheduler(input.num_inference_steps.unwrap_or(50).into()); 74 | let text_embeddings = self.tokenize(input.prompt)?; 75 | 76 | let mut outputs = Vec::new(); 77 | for idx in 0..input.num_outputs.unwrap_or(1) { 78 | if let Some(seed) = input.seed { 79 | tch::manual_seed((seed + idx as u32).into()); 80 | } else { 81 | tch::manual_seed(-1); 82 | } 83 | 84 | let mut latents = Tensor::randn( 85 | [1, 4, self.sd_config.height / 8, self.sd_config.width / 8], 86 | (Kind::Float, self.devices.get("unet")), 87 | ); 88 | 89 | latents *= scheduler.init_noise_sigma(); 90 | 91 | for ×tep in scheduler.timesteps().iter() { 92 | let latent_model_input = Tensor::cat(&[&latents, &latents], 0); 93 | 94 | let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep); 95 | let noise_pred = 96 | self.unet 97 | .forward(&latent_model_input, timestep as f64, &text_embeddings); 98 | let noise_pred = noise_pred.chunk(2, 0); 99 | let (noise_pred_uncond, noise_pred_text) = (&noise_pred[0], &noise_pred[1]); 100 | let noise_pred = noise_pred_uncond 101 | + (noise_pred_text - noise_pred_uncond) * input.guidance_scale.unwrap_or(7.5); 102 | latents = scheduler.step(&noise_pred, timestep, &latents); 103 | } 104 | 105 | let latents = latents.to(self.devices.get("vae")); 106 | let image = self.vae.decode(&(&latents / 0.18215)); 107 | let image = (image / 2 + 0.5).clamp(0., 1.).to_device(Device::Cpu); 108 | let image = (image * 255.).to_kind(Kind::Uint8); 109 | 110 | let final_image = PathBuf::from(format!("output-{idx}.png")); 111 | tch::vision::image::save(&image, &final_image)?; 112 | outputs.push(final_image.into()); 113 | } 114 | 115 | Ok(outputs) 116 | } 117 | } 118 | 119 | impl StableDiffusion { 120 | fn tokenize(&self, prompt: String) -> Result { 121 | let tokens = self.tokenizer.encode(&prompt)?; 122 | let tokens: Vec = tokens.into_iter().map(|x| x as i64).collect(); 123 | let tokens = Tensor::from_slice(&tokens) 124 | .view((1, -1)) 125 | .to(self.devices.get("clip")); 126 | 127 | let uncond_tokens = self.tokenizer.encode("")?; 128 | let uncond_tokens: Vec = uncond_tokens.into_iter().map(|x| x as i64).collect(); 129 | let uncond_tokens = Tensor::from_slice(&uncond_tokens) 130 | .view((1, -1)) 131 | .to(self.devices.get("clip")); 132 | 133 | let text_embeddings = self.text_model.forward(&tokens); 134 | let uncond_embeddings = self.text_model.forward(&uncond_tokens); 135 | 136 | Ok(Tensor::cat(&[uncond_embeddings, text_embeddings], 0).to(self.devices.get("unet"))) 137 | } 138 | } 139 | 140 | cog_rust::start!(StableDiffusion); 141 | -------------------------------------------------------------------------------- /examples/stable-diffusion/weights/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/m1guelpf/cog-rust/4d6666a2245a6ae9a6e14973f3a430a94e3fd4f9/examples/stable-diffusion/weights/.gitkeep -------------------------------------------------------------------------------- /lib/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "cog-rust" 3 | version = "1.0.14" 4 | description = "Rust containers for machine learning" 5 | readme = { workspace = true } 6 | edition = { workspace = true } 7 | authors = { workspace = true } 8 | license = { workspace = true } 9 | keywords = { workspace = true } 10 | categories = { workspace = true } 11 | repository = { workspace = true } 12 | 13 | [dependencies] 14 | anyhow = "1.0.71" 15 | flume = "0.10.14" 16 | serde = "1.0.163" 17 | base64 = "0.21.2" 18 | tracing = "0.1.37" 19 | indexmap = "1.9.3" 20 | titlecase = "2.2.1" 21 | map-macro = "0.2.6" 22 | itertools = "0.11.0" 23 | thiserror = "1.0.40" 24 | mime_guess = "2.0.4" 25 | lazy_static = "1.4.0" 26 | atomic_enum = "0.2.0" 27 | serde_json = "1.0.96" 28 | async-trait = "0.1.68" 29 | percent-encoding = "2.3.0" 30 | uuid = { version = "1.3.3", features = ["v4"] } 31 | url = { version = "2.4.0", features = ["serde"] } 32 | cog-core = { path = "../core", version = "0.2.0" } 33 | clap = { version = "4.3.21", features = ["derive"] } 34 | axum = { version = "0.6.18", features = ["headers"] } 35 | tokio = { version = "1.28.2", features = ["full"] } 36 | chrono = { version = "0.4.26", features = ["serde"] } 37 | axum-jsonschema = { version = "0.6.0", features = ["aide"] } 38 | jsonschema = { version = "0.17.0", default-features = false } 39 | schemars = { version = "0.8.12", features = ["chrono", "url"] } 40 | reqwest = { version = "0.11.18", features = ["json", "blocking"] } 41 | aide = { version = "0.11.0", features = ["axum", "axum-headers"] } 42 | tracing-subscriber = { version = "0.3.17", features = ["env-filter"] } 43 | tree_magic_mini = { version = "3.0.3", features = [ 44 | "tree_magic_db", 45 | "with-gpl-data", 46 | ] } 47 | -------------------------------------------------------------------------------- /lib/src/errors.rs: -------------------------------------------------------------------------------- 1 | use aide::OperationOutput; 2 | use axum::{ 3 | http::StatusCode, 4 | response::{IntoResponse, Response}, 5 | Json, 6 | }; 7 | use cog_core::http::ValidationError; 8 | use jsonschema::ErrorIterator; 9 | use serde_json::{json, Value}; 10 | 11 | use crate::prediction::Error as PredictionError; 12 | 13 | #[derive(Debug)] 14 | pub struct HTTPError { 15 | detail: Value, 16 | status_code: StatusCode, 17 | } 18 | 19 | impl HTTPError { 20 | pub fn new(detail: &str) -> Self { 21 | Self { 22 | detail: detail.into(), 23 | status_code: StatusCode::UNPROCESSABLE_ENTITY, 24 | } 25 | } 26 | 27 | pub const fn with_status(mut self, status_code: StatusCode) -> Self { 28 | self.status_code = status_code; 29 | self 30 | } 31 | } 32 | 33 | impl IntoResponse for HTTPError { 34 | fn into_response(self) -> Response { 35 | (self.status_code, Json(json!({ "detail": self.detail }))).into_response() 36 | } 37 | } 38 | 39 | impl OperationOutput for HTTPError { 40 | type Inner = Self; 41 | } 42 | 43 | #[derive(Debug, Clone, thiserror::Error, serde::Deserialize, serde::Serialize)] 44 | #[error("Validation Errors")] 45 | pub struct ValidationErrorSet { 46 | pub errors: Vec, 47 | } 48 | 49 | impl ValidationErrorSet { 50 | pub fn fill_loc(mut self, loc: &[&str]) -> Self { 51 | self.errors 52 | .iter_mut() 53 | .map(|error| { 54 | error.loc = loc 55 | .iter() 56 | .map(ToString::to_string) 57 | .chain(error.loc.clone()) 58 | .collect(); 59 | }) 60 | .for_each(drop); 61 | 62 | self 63 | } 64 | } 65 | 66 | impl From> for ValidationErrorSet { 67 | fn from(e: ErrorIterator<'_>) -> Self { 68 | Self { 69 | errors: e 70 | .map(|e| ValidationError { 71 | msg: e.to_string(), 72 | loc: e.instance_path.into_vec(), 73 | }) 74 | .collect(), 75 | } 76 | } 77 | } 78 | 79 | #[allow(clippy::fallible_impl_from)] 80 | impl From for HTTPError { 81 | fn from(e: ValidationErrorSet) -> Self { 82 | Self { 83 | status_code: StatusCode::UNPROCESSABLE_ENTITY, 84 | detail: serde_json::to_value(e.errors).unwrap(), 85 | } 86 | } 87 | } 88 | 89 | #[allow(clippy::fallible_impl_from)] 90 | impl From for HTTPError { 91 | fn from(e: PredictionError) -> Self { 92 | match e { 93 | PredictionError::Unknown => Self { 94 | status_code: StatusCode::NOT_FOUND, 95 | detail: serde_json::to_value(e.to_string()).unwrap(), 96 | }, 97 | PredictionError::AlreadyRunning => Self { 98 | status_code: StatusCode::CONFLICT, 99 | detail: serde_json::to_value(e.to_string()).unwrap(), 100 | }, 101 | PredictionError::Validation(e) => e.into(), 102 | PredictionError::NotComplete | PredictionError::Receiver(_) => Self { 103 | status_code: StatusCode::INTERNAL_SERVER_ERROR, 104 | detail: serde_json::to_value(e.to_string()).unwrap(), 105 | }, 106 | } 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /lib/src/helpers/headers.rs: -------------------------------------------------------------------------------- 1 | use axum::{ 2 | headers::{Error, Header}, 3 | http::{HeaderName, HeaderValue}, 4 | }; 5 | use itertools::Itertools; 6 | use lazy_static::lazy_static; 7 | use percent_encoding::{percent_decode_str, percent_encode, NON_ALPHANUMERIC}; 8 | use std::{borrow::Cow, collections::HashMap}; 9 | 10 | lazy_static! { 11 | static ref PREFER: HeaderName = HeaderName::from_lowercase(b"prefer").unwrap(); 12 | } 13 | 14 | #[derive(Debug, Clone, PartialEq, Eq, Default)] 15 | pub struct Prefer(pub HashMap); 16 | 17 | impl Prefer { 18 | pub fn has(&self, key: &str) -> bool { 19 | self.0.contains_key(key) 20 | } 21 | } 22 | 23 | impl Header for Prefer { 24 | fn name() -> &'static HeaderName { 25 | &PREFER 26 | } 27 | 28 | fn decode<'i, I>(values: &mut I) -> Result 29 | where 30 | Self: Sized, 31 | I: Iterator, 32 | { 33 | let value = values.next().ok_or_else(Error::invalid)?; 34 | 35 | let preferences = value 36 | .to_str() 37 | .map_err(|_| Error::invalid())? 38 | .split(',') 39 | .map(str::trim) 40 | .map(|s| { 41 | let mut split = s.splitn(2, '='); 42 | let (key, value) = (split.next().unwrap(), split.next().unwrap_or_default()); 43 | 44 | ( 45 | key.to_string(), 46 | percent_decode_str(value) 47 | .decode_utf8() 48 | .unwrap_or(Cow::Borrowed(value)) 49 | .to_string(), 50 | ) 51 | }) 52 | .collect::>(); 53 | 54 | Ok(Self(preferences)) 55 | } 56 | 57 | fn encode>(&self, values: &mut E) { 58 | let value = self 59 | .0 60 | .iter() 61 | .sorted() 62 | .map(|(key, value)| { 63 | format!( 64 | "{key}{}{}", 65 | if value.is_empty() { "" } else { "=" }, 66 | percent_encode(value.as_bytes(), NON_ALPHANUMERIC) 67 | ) 68 | }) 69 | .collect::>() 70 | .join(","); 71 | 72 | let value = HeaderValue::from_bytes(value.as_bytes()).unwrap(); 73 | values.extend(std::iter::once(value)); 74 | } 75 | } 76 | 77 | #[cfg(test)] 78 | mod tests { 79 | use super::*; 80 | use axum::http::HeaderMap; 81 | use map_macro::hash_map; 82 | 83 | #[test] 84 | fn header_is_parsed_correctly() { 85 | let mut headers = HeaderMap::new(); 86 | headers.insert( 87 | "Prefer", 88 | HeaderValue::from_static("wait=10, timeout=5, respond-async"), 89 | ); 90 | 91 | let prefer = Prefer::decode(&mut headers.get_all("Prefer").iter()).unwrap(); 92 | 93 | assert_eq!( 94 | prefer, 95 | Prefer(hash_map! { 96 | "wait".to_string() => "10".to_string(), 97 | "timeout".to_string() => "5".to_string(), 98 | "respond-async".to_string() => String::new(), 99 | }) 100 | ); 101 | } 102 | 103 | #[test] 104 | fn header_is_encoded_correctly() { 105 | let prefer = Prefer(hash_map! { 106 | "wait".to_string() => "10".to_string(), 107 | "timeout".to_string() => "5".to_string(), 108 | "respond-async".to_string() => String::new(), 109 | }); 110 | 111 | let mut values = Vec::new(); 112 | prefer.encode(&mut values); 113 | 114 | assert_eq!( 115 | values, 116 | vec![HeaderValue::from_static("respond-async,timeout=5,wait=10")] 117 | ); 118 | } 119 | 120 | #[test] 121 | fn has_returns_true_if_key_exists() { 122 | let prefer = Prefer(hash_map! { 123 | "wait".to_string() => "10".to_string(), 124 | "timeout".to_string() => "5".to_string(), 125 | "respond-async".to_string() => String::new(), 126 | }); 127 | 128 | assert!(prefer.has("wait")); 129 | assert!(prefer.has("timeout")); 130 | assert!(prefer.has("respond-async")); 131 | assert!(!prefer.has("foo")); 132 | } 133 | } 134 | -------------------------------------------------------------------------------- /lib/src/helpers/mod.rs: -------------------------------------------------------------------------------- 1 | use base64::{engine::general_purpose::STANDARD as Base64, DecodeError, Engine}; 2 | use url::Url; 3 | 4 | pub mod headers; 5 | pub mod openapi; 6 | 7 | pub fn base64_encode>(bytes: T) -> String { 8 | Base64.encode(bytes) 9 | } 10 | 11 | pub fn base64_decode>(bytes: T) -> Result, DecodeError> { 12 | Base64.decode(bytes) 13 | } 14 | 15 | /// Append a path to a URL. 16 | /// This is a workaround for the fact that `Url::join` will get rid of the last path segment if it doesn't end with a slash. 17 | pub fn url_join(url: &Url, path: &str) -> Url { 18 | let mut url = url.clone(); 19 | let mut path_parts = url.path_segments_mut().unwrap(); 20 | path_parts.pop_if_empty(); 21 | 22 | path_parts.push(path); 23 | drop(path_parts); 24 | 25 | url 26 | } 27 | -------------------------------------------------------------------------------- /lib/src/helpers/openapi.rs: -------------------------------------------------------------------------------- 1 | use aide::openapi::{MediaType, OpenApi, ReferenceOr, Response, SchemaObject, StatusCode}; 2 | use axum::http::Method; 3 | use schemars::{gen::SchemaGenerator, schema::Schema, JsonSchema}; 4 | 5 | pub fn schema_with_properties( 6 | generator: &mut SchemaGenerator, 7 | cb: impl Fn(String, &mut schemars::schema::SchemaObject, usize), 8 | ) -> Schema { 9 | let mut schema = generator.root_schema_for::().schema; 10 | let metadata = schema.metadata(); 11 | 12 | metadata.title = Some(metadata.title.as_ref().map_or_else( 13 | || T::schema_name(), 14 | |title| title.split('_').next().unwrap().to_string(), 15 | )); 16 | 17 | let object = schema.object(); 18 | for (index, (name, property)) in object.properties.clone().into_iter().enumerate() { 19 | let mut property: schemars::schema::SchemaObject = property.clone().into_object(); 20 | 21 | cb(name.clone(), &mut property, index); 22 | object.properties.insert(name, property.into()); 23 | } 24 | 25 | schemars::schema::Schema::Object(schema) 26 | } 27 | 28 | pub fn replace_request_schema( 29 | api: &mut OpenApi, 30 | path: &str, 31 | (method, media_type): (Method, &str), 32 | schema: schemars::schema::SchemaObject, 33 | ) -> Option<()> { 34 | let paths = api.paths.as_mut()?; 35 | let item = paths.paths.get_mut(path)?.as_item_mut()?; 36 | let operation = match method { 37 | Method::GET => item.get.as_mut()?, 38 | Method::PUT => item.put.as_mut()?, 39 | Method::POST => item.post.as_mut()?, 40 | Method::HEAD => item.head.as_mut()?, 41 | Method::TRACE => item.trace.as_mut()?, 42 | Method::DELETE => item.delete.as_mut()?, 43 | Method::OPTIONS => item.options.as_mut()?, 44 | _ => return None, 45 | }; 46 | 47 | let body = operation.request_body.as_mut()?.as_item_mut()?; 48 | 49 | body.content.get_mut(media_type)?.schema = Some(SchemaObject { 50 | example: None, 51 | external_docs: None, 52 | json_schema: Schema::Object(schema), 53 | }); 54 | 55 | Some(()) 56 | } 57 | 58 | pub fn replace_response_schema( 59 | api: &mut OpenApi, 60 | path: &str, 61 | (method, status_code, media_type): (Method, StatusCode, String), 62 | json_schema: schemars::schema::SchemaObject, 63 | ) -> Option<()> { 64 | let paths = api.paths.as_mut()?; 65 | let item = paths.paths.get_mut(path)?.as_item_mut()?; 66 | let operation = match method { 67 | Method::GET => item.get.as_mut()?, 68 | Method::PUT => item.put.as_mut()?, 69 | Method::POST => item.post.as_mut()?, 70 | Method::HEAD => item.head.as_mut()?, 71 | Method::TRACE => item.trace.as_mut()?, 72 | Method::DELETE => item.delete.as_mut()?, 73 | Method::OPTIONS => item.options.as_mut()?, 74 | _ => return None, 75 | }; 76 | 77 | let mut responses = operation.responses.clone().unwrap_or_default(); 78 | let response = responses.responses.get(&status_code).cloned(); 79 | let mut response = response.unwrap_or_else(|| ReferenceOr::Item(Response::default())); 80 | let response = response.as_item_mut()?; 81 | 82 | response.content.insert( 83 | media_type, 84 | MediaType { 85 | schema: Some(SchemaObject { 86 | example: None, 87 | external_docs: None, 88 | json_schema: Schema::Object(json_schema), 89 | }), 90 | ..Default::default() 91 | }, 92 | ); 93 | 94 | responses 95 | .responses 96 | .insert(status_code, ReferenceOr::Item(response.clone())); 97 | operation.responses = Some(responses); 98 | 99 | Some(()) 100 | } 101 | -------------------------------------------------------------------------------- /lib/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![warn(clippy::all, clippy::pedantic, clippy::nursery)] 2 | 3 | use anyhow::Result; 4 | use clap::Parser; 5 | use tracing_subscriber::{ 6 | prelude::__tracing_subscriber_SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer, 7 | }; 8 | 9 | pub use cog_core::{Cog, CogResponse}; 10 | pub use spec::Path; 11 | 12 | mod errors; 13 | mod helpers; 14 | mod prediction; 15 | mod routes; 16 | mod runner; 17 | mod server; 18 | mod shutdown; 19 | mod spec; 20 | mod webhooks; 21 | 22 | #[derive(Debug, clap::Parser)] 23 | pub(crate) struct Cli { 24 | /// Dump the schema and exit 25 | #[clap(long)] 26 | dump_schema_and_exit: bool, 27 | 28 | /// Ignore SIGTERM and wait for a request to /shutdown (or a SIGINT) before exiting 29 | #[arg(long, default_missing_value = "true", require_equals = true, num_args=0..=1, action = clap::ArgAction::Set)] 30 | await_explicit_shutdown: Option, 31 | 32 | /// An endpoint for Cog to PUT output files to 33 | #[clap(long)] 34 | upload_url: Option, 35 | } 36 | 37 | /// Start the server with the given model. 38 | /// 39 | /// # Errors 40 | /// 41 | /// This function will return an error if the PORT environment variable is set but cannot be parsed, or if the server fails to start. 42 | pub async fn start() -> Result<()> { 43 | let args = Cli::parse(); 44 | 45 | if !args.dump_schema_and_exit { 46 | tracing_subscriber::registry() 47 | .with(tracing_subscriber::fmt::layer().with_filter( 48 | EnvFilter::try_from_default_env().unwrap_or_else(|_| "cog_rust=info".into()), 49 | )) 50 | .init(); 51 | } 52 | 53 | server::start::(args).await 54 | } 55 | 56 | #[macro_export] 57 | /// Start the server with the given model. 58 | macro_rules! start { 59 | ($struct_name:ident) => { 60 | #[tokio::main] 61 | async fn main() { 62 | cog_rust::start::<$struct_name>().await.unwrap(); 63 | } 64 | }; 65 | } 66 | -------------------------------------------------------------------------------- /lib/src/prediction.rs: -------------------------------------------------------------------------------- 1 | use chrono::{DateTime, Utc}; 2 | use cog_core::http::{Request, Response, Status}; 3 | use map_macro::hash_map; 4 | use serde_json::Value; 5 | use std::{ 6 | future::Future, 7 | sync::{atomic::Ordering, Arc}, 8 | time::Duration, 9 | }; 10 | use tokio::sync::RwLock; 11 | 12 | use crate::{ 13 | errors::ValidationErrorSet, 14 | runner::{Error as RunnerError, Health, Runner, RUNNER_HEALTH}, 15 | shutdown::Shutdown, 16 | webhooks::WebhookSender, 17 | Cog, 18 | }; 19 | 20 | pub type Extension = axum::Extension>>; 21 | 22 | #[derive(Debug, Clone, thiserror::Error)] 23 | pub enum Error { 24 | #[error("Attempted to re-initialize a prediction")] 25 | AlreadyRunning, 26 | 27 | #[error("Prediction is not yet complete")] 28 | NotComplete, 29 | 30 | #[error("The requested prediction does not exist")] 31 | Unknown, 32 | 33 | #[error("Failed to wait for prediction: {0}")] 34 | Receiver(#[from] flume::RecvError), 35 | 36 | #[error("Failed to run prediction: {0}")] 37 | Validation(#[from] ValidationErrorSet), 38 | } 39 | 40 | pub struct Prediction { 41 | runner: Runner, 42 | pub status: Status, 43 | pub id: Option, 44 | pub shutdown: Shutdown, 45 | webhooks: WebhookSender, 46 | cancel: flume::Sender<()>, 47 | pub request: Option, 48 | pub response: Option, 49 | complete: Option>, 50 | } 51 | 52 | impl Prediction { 53 | pub fn setup(shutdown: Shutdown) -> Self { 54 | let (cancel_tx, cancel_rx) = flume::unbounded(); 55 | 56 | Self { 57 | id: None, 58 | request: None, 59 | complete: None, 60 | response: None, 61 | cancel: cancel_tx, 62 | status: Status::Idle, 63 | shutdown: shutdown.clone(), 64 | webhooks: WebhookSender::new().unwrap(), 65 | runner: Runner::new::(shutdown, cancel_rx), 66 | } 67 | } 68 | 69 | pub fn init(&mut self, id: Option, req: Request) -> Result<&mut Self, Error> { 70 | if !matches!(self.status, Status::Idle) { 71 | tracing::debug!("Attempted to re-initialize a prediction"); 72 | return Err(Error::AlreadyRunning); 73 | } 74 | 75 | self.validate(&req.input) 76 | .map_err(|e| e.fill_loc(&["body", "input"]))?; 77 | 78 | tracing::debug!("Initializing prediction: {id:?}"); 79 | 80 | self.id = id; 81 | self.request = Some(req); 82 | self.status = Status::Starting; 83 | 84 | Ok(self) 85 | } 86 | 87 | pub fn validate(&self, input: &Value) -> Result<(), ValidationErrorSet> { 88 | self.runner.validate(input) 89 | } 90 | 91 | pub async fn run(&mut self) -> Result { 92 | self.process()?.await; 93 | 94 | self.result() 95 | } 96 | 97 | pub async fn wait_for(&self, id: String) -> Result { 98 | if self.id != Some(id.clone()) { 99 | tracing::debug!("Attempted to wait for prediction with unknown ID: {id:?}"); 100 | return Err(Error::Unknown); 101 | } 102 | 103 | if let Some(response) = self.response.clone() { 104 | tracing::debug!("Prediction already complete: {id:?}"); 105 | return Ok(response); 106 | } 107 | 108 | if !matches!(self.status, Status::Processing) { 109 | tracing::debug!("Attempted to wait for prediction that isn't running: {id:?}"); 110 | return Err(Error::AlreadyRunning); 111 | } 112 | 113 | tracing::debug!("Waiting for prediction: {id:?}"); 114 | let complete = self.complete.as_ref().unwrap(); 115 | Ok(complete.recv_async().await?) 116 | } 117 | 118 | pub fn process(&mut self) -> Result + '_, Error> { 119 | if !matches!(self.status, Status::Starting) { 120 | tracing::debug!( 121 | "Attempted to process prediction while not ready: {:?}", 122 | self.id 123 | ); 124 | return Err(Error::AlreadyRunning); 125 | } 126 | 127 | let req = self.request.clone().unwrap(); 128 | self.status = Status::Processing; 129 | 130 | let (complete_tx, complete_rx) = flume::bounded(1); 131 | self.complete = Some(complete_rx); 132 | 133 | Ok(async move { 134 | let started_at = Utc::now(); 135 | tracing::debug!("Running prediction: {:?}", self.id); 136 | 137 | self.status = Status::Processing; 138 | self.response = Some(Response::starting(self.id.clone(), req.clone())); 139 | if let Err(e) = self.webhooks.starting(self).await { 140 | tracing::error!("Failed to send start webhook for prediction: {e:?}",); 141 | }; 142 | 143 | tokio::select! { 144 | () = self.shutdown.handle() => { 145 | tracing::debug!("Shutdown requested. Cancelling running prediction: {:?}", self.id); 146 | return; 147 | }, 148 | output = self.runner.run(req.clone()) => { 149 | tracing::debug!("Prediction complete: {:?}", self.id); 150 | 151 | match output { 152 | Ok((output, predict_time)) => { 153 | self.status = Status::Succeeded; 154 | self.response = Some(Response::success(self.id.clone(), req, output, predict_time, started_at)); 155 | }, 156 | Err(RunnerError::Canceled) => { 157 | self.status = Status::Canceled; 158 | self.response = Some(Response::canceled(self.id.clone(), req, started_at)); 159 | 160 | }, 161 | Err(error) => { 162 | self.status = Status::Failed; 163 | self.response = Some(Response::error(self.id.clone(), req, &error, started_at)); 164 | } 165 | } 166 | 167 | if let Err(e) = self.webhooks.finished(self, self.response.clone().unwrap()).await { 168 | tracing::error!("Failed to send finished webhook for prediction: {e:?}",); 169 | }; 170 | } 171 | } 172 | complete_tx.send(self.response.clone().unwrap()).unwrap(); 173 | }) 174 | } 175 | 176 | pub fn result(&mut self) -> Result { 177 | if !matches!( 178 | self.status, 179 | Status::Succeeded | Status::Failed | Status::Canceled 180 | ) { 181 | tracing::debug!( 182 | "Attempted to get result of prediction that is not complete: {:?}", 183 | self.id 184 | ); 185 | return Err(Error::NotComplete); 186 | } 187 | 188 | tracing::debug!("Getting result of prediction: {:?}", self.id); 189 | let response = self.response.clone().ok_or(Error::NotComplete)?; 190 | self.reset(); 191 | 192 | Ok(response) 193 | } 194 | 195 | pub fn cancel(&mut self, id: &str) -> Result<&mut Self, Error> { 196 | if self.id != Some(id.to_string()) { 197 | tracing::debug!("Attempted to cancel prediction with unknown ID: {id}"); 198 | return Err(Error::Unknown); 199 | } 200 | 201 | if !matches!(self.status, Status::Processing) { 202 | tracing::debug!("Attempted to cancel prediction that is not running: {id}"); 203 | return Err(Error::AlreadyRunning); 204 | } 205 | 206 | tracing::debug!("Canceling prediction: {id}"); 207 | self.cancel.send(()).unwrap(); 208 | self.status = Status::Canceled; 209 | 210 | Ok(self) 211 | } 212 | 213 | pub fn reset(&mut self) { 214 | tracing::debug!("Resetting prediction"); 215 | 216 | self.id = None; 217 | self.request = None; 218 | self.response = None; 219 | self.complete = None; 220 | self.status = Status::Idle; 221 | } 222 | 223 | pub fn extension(self) -> Extension { 224 | axum::Extension(Arc::new(RwLock::new(self))) 225 | } 226 | } 227 | 228 | pub struct SyncGuard<'a> { 229 | prediction: tokio::sync::RwLockWriteGuard<'a, Prediction>, 230 | } 231 | 232 | impl<'a> SyncGuard<'a> { 233 | pub fn new(prediction: tokio::sync::RwLockWriteGuard<'a, Prediction>) -> Self { 234 | Self { prediction } 235 | } 236 | 237 | pub fn init(&mut self, id: Option, req: Request) -> Result<&mut Self, Error> { 238 | self.prediction.init(id, req)?; 239 | Ok(self) 240 | } 241 | 242 | pub async fn run(&mut self) -> Result { 243 | self.prediction.run().await 244 | } 245 | } 246 | 247 | impl Drop for SyncGuard<'_> { 248 | fn drop(&mut self) { 249 | tracing::debug!("SyncGuard dropped, resetting prediction"); 250 | 251 | self.prediction.reset(); 252 | if matches!(RUNNER_HEALTH.load(Ordering::SeqCst), Health::Busy) { 253 | self.prediction.cancel.send(()).unwrap(); 254 | } 255 | } 256 | } 257 | 258 | pub trait ResponseHelpers { 259 | fn starting(id: Option, req: Request) -> Self; 260 | fn success( 261 | id: Option, 262 | req: Request, 263 | output: Value, 264 | predict_time: Duration, 265 | started_at: DateTime, 266 | ) -> Self; 267 | fn error( 268 | id: Option, 269 | req: Request, 270 | error: &RunnerError, 271 | started_at: DateTime, 272 | ) -> Self; 273 | fn canceled(id: Option, req: Request, started_at: DateTime) -> Self; 274 | } 275 | 276 | impl ResponseHelpers for Response { 277 | fn success( 278 | id: Option, 279 | req: Request, 280 | output: Value, 281 | predict_time: Duration, 282 | started_at: DateTime, 283 | ) -> Self { 284 | Self { 285 | id, 286 | output: Some(output), 287 | input: Some(req.input), 288 | status: Status::Succeeded, 289 | started_at: Some(started_at), 290 | completed_at: Some(Utc::now()), 291 | metrics: Some(hash_map! { 292 | "predict_time".to_string() => predict_time.as_secs_f64().into() 293 | }), 294 | ..Self::default() 295 | } 296 | } 297 | fn error( 298 | id: Option, 299 | req: Request, 300 | error: &RunnerError, 301 | started_at: DateTime, 302 | ) -> Self { 303 | Self { 304 | id, 305 | input: Some(req.input), 306 | status: Status::Failed, 307 | started_at: Some(started_at), 308 | error: Some(error.to_string()), 309 | ..Self::default() 310 | } 311 | } 312 | 313 | fn starting(id: Option, req: Request) -> Self { 314 | Self { 315 | id, 316 | input: Some(req.input), 317 | status: Status::Processing, 318 | started_at: Some(Utc::now()), 319 | ..Self::default() 320 | } 321 | } 322 | 323 | fn canceled(id: Option, req: Request, started_at: DateTime) -> Self { 324 | Self { 325 | id, 326 | input: Some(req.input), 327 | status: Status::Canceled, 328 | started_at: Some(started_at), 329 | ..Self::default() 330 | } 331 | } 332 | } 333 | -------------------------------------------------------------------------------- /lib/src/routes/docs.rs: -------------------------------------------------------------------------------- 1 | use aide::{ 2 | axum::{routing::get, ApiRouter}, 3 | openapi::OpenApi, 4 | }; 5 | use axum::{response::Html, Extension}; 6 | use axum_jsonschema::Json; 7 | 8 | pub fn handler() -> ApiRouter { 9 | ApiRouter::new() 10 | .route("/docs", get(swagger)) 11 | .route("/openapi.json", get(openapi_schema)) 12 | } 13 | 14 | #[allow(clippy::unused_async)] 15 | async fn openapi_schema(Extension(openapi): Extension) -> Json { 16 | Json(openapi) 17 | } 18 | 19 | #[allow(clippy::unused_async)] 20 | async fn swagger() -> Html { 21 | Html(SWAGGER_UI_TEMPLATE.replace("{:spec_url}", "/openapi.json")) 22 | } 23 | 24 | const SWAGGER_UI_TEMPLATE: &str = r#" 25 | 26 | 27 | 28 | 29 | Cog Docs 30 | 31 | 32 | 33 | 34 |
35 | 36 | 37 | 42 | 43 | 44 | "#; 45 | -------------------------------------------------------------------------------- /lib/src/routes/mod.rs: -------------------------------------------------------------------------------- 1 | use aide::axum::ApiRouter; 2 | 3 | mod docs; 4 | mod predict; 5 | mod system; 6 | 7 | pub fn handler() -> ApiRouter { 8 | ApiRouter::new() 9 | .merge(system::handler()) 10 | .merge(predict::handler()) 11 | .merge(docs::handler()) 12 | } 13 | -------------------------------------------------------------------------------- /lib/src/routes/predict.rs: -------------------------------------------------------------------------------- 1 | use aide::axum::{ 2 | routing::{post, put}, 3 | ApiRouter, 4 | }; 5 | use axum::{extract::Path, http::StatusCode, Extension, TypedHeader}; 6 | use axum_jsonschema::Json; 7 | use cog_core::http::Status; 8 | 9 | use crate::{ 10 | errors::HTTPError, 11 | helpers::headers::Prefer, 12 | prediction::{Extension as ExtractPrediction, ResponseHelpers, SyncGuard}, 13 | }; 14 | 15 | pub fn handler() -> ApiRouter { 16 | ApiRouter::new() 17 | .api_route("/predictions", post(create_prediction)) 18 | .api_route("/predictions/:prediction_id", put(create_prediction)) 19 | .api_route( 20 | "/predictions/:prediction_id/cancel", 21 | post(cancel_prediction), 22 | ) 23 | } 24 | 25 | async fn create_prediction( 26 | id: Option>, 27 | prefer: Option>, 28 | Extension(prediction): ExtractPrediction, 29 | Json(req): Json, 30 | ) -> Result<(StatusCode, Json), HTTPError> { 31 | let id = id.map(|id| id.0); 32 | let respond_async = prefer 33 | .map(|prefer| prefer.0) 34 | .unwrap_or_default() 35 | .has("respond-async"); 36 | 37 | tracing::debug!( 38 | "Received {}prediction request{}.", 39 | if respond_async { "async " } else { "" }, 40 | id.as_ref() 41 | .map_or(String::new(), |id| format!(" with id {id}")), 42 | ); 43 | tracing::trace!("{req:?}"); 44 | 45 | let r_prediction = prediction.read().await; 46 | 47 | // If a named prediction is already running... 48 | if let Some(prediction_id) = r_prediction.id.clone() { 49 | // ...and the request is for a different prediction, return an error. 50 | if let Some(id) = id.clone() { 51 | if Some(prediction_id.clone()) != Some(id.clone()) { 52 | tracing::debug!( 53 | "Trying to run a named prediction {id} while another prediction {prediction_id} is running" 54 | ); 55 | return Err(HTTPError::new("Already running a prediction") 56 | .with_status(StatusCode::CONFLICT)); 57 | } 58 | 59 | // ...and this is an async request, return the current response. 60 | if respond_async { 61 | return Ok(( 62 | StatusCode::ACCEPTED, 63 | Json(r_prediction.response.clone().unwrap()), 64 | )); 65 | } 66 | 67 | // wait for the current prediction to complete 68 | return Ok((StatusCode::OK, Json(r_prediction.wait_for(id).await?))); 69 | } 70 | } 71 | 72 | // If the request is synchronous, run the prediction and return the result. 73 | if !respond_async { 74 | drop(r_prediction); 75 | let mut prediction = SyncGuard::new(prediction.write().await); 76 | 77 | return Ok((StatusCode::OK, Json(prediction.init(id, req)?.run().await?))); 78 | } 79 | 80 | // If there's a running prediction, return an error. 81 | if !matches!(r_prediction.status, Status::Idle) { 82 | return Err( 83 | HTTPError::new("Already running a prediction").with_status(StatusCode::CONFLICT) 84 | ); 85 | } 86 | 87 | // Throw an error if the request is invalid. 88 | r_prediction.validate(&req.input)?; 89 | drop(r_prediction); 90 | 91 | let thread_req = req.clone(); 92 | let thread_id = id.clone(); 93 | tokio::spawn(async move { 94 | tracing::debug!("Running prediction asynchronously: {:?}", thread_id); 95 | 96 | let mut prediction = prediction.write().await; 97 | 98 | prediction.init(thread_id.clone(), thread_req).unwrap(); 99 | prediction.process().unwrap().await; 100 | prediction.reset(); 101 | drop(prediction); 102 | 103 | tracing::debug!("Asynchronous prediction complete: {thread_id:?}"); 104 | }); 105 | 106 | Ok(( 107 | StatusCode::ACCEPTED, 108 | Json(cog_core::http::Response::starting(id, req)), 109 | )) 110 | } 111 | 112 | async fn cancel_prediction( 113 | Path(id): Path, 114 | Extension(prediction): ExtractPrediction, 115 | ) -> Result, HTTPError> { 116 | let mut prediction = prediction.write().await; 117 | prediction.cancel(&id)?; 118 | drop(prediction); 119 | 120 | Ok(Json(())) 121 | } 122 | -------------------------------------------------------------------------------- /lib/src/routes/system.rs: -------------------------------------------------------------------------------- 1 | use std::sync::atomic::Ordering; 2 | 3 | use aide::axum::{ 4 | routing::{get, post}, 5 | ApiRouter, 6 | }; 7 | use axum::Extension; 8 | use axum_jsonschema::Json; 9 | use chrono::Utc; 10 | use cog_core::http::Status; 11 | use schemars::JsonSchema; 12 | 13 | use crate::{ 14 | runner::{Health, RUNNER_HEALTH}, 15 | shutdown::Agent as Shutdown, 16 | }; 17 | 18 | pub fn handler() -> ApiRouter { 19 | ApiRouter::new() 20 | .api_route("/", get(root)) 21 | .api_route("/health-check", get(health_check)) 22 | .api_route("/shutdown", post(shutdown)) 23 | } 24 | 25 | #[derive(Debug, serde::Serialize, JsonSchema)] 26 | pub struct RootResponse { 27 | /// Relative URL to Swagger UI 28 | pub docs_url: String, 29 | /// Relative URL to OpenAPI specification 30 | pub openapi_url: String, 31 | } 32 | 33 | #[allow(clippy::unused_async)] 34 | pub async fn root() -> Json { 35 | Json(RootResponse { 36 | docs_url: "/docs".to_string(), 37 | openapi_url: "/openapi.json".to_string(), 38 | }) 39 | } 40 | 41 | #[derive(Debug, serde::Serialize, JsonSchema)] 42 | pub struct HealthCheckSetup { 43 | /// Setup logs 44 | pub logs: String, 45 | /// Setup status 46 | pub status: Status, 47 | /// Setup started time 48 | pub started_at: String, 49 | /// Setup completed time 50 | pub completed_at: String, 51 | } 52 | 53 | #[derive(Debug, serde::Serialize, JsonSchema)] 54 | pub struct HealthCheck { 55 | /// Current health status 56 | pub status: Health, 57 | /// Setup information 58 | pub setup: HealthCheckSetup, 59 | } 60 | 61 | #[allow(clippy::unused_async)] 62 | pub async fn health_check() -> Json { 63 | let status = RUNNER_HEALTH.load(Ordering::SeqCst); 64 | 65 | Json(HealthCheck { 66 | status, 67 | setup: HealthCheckSetup { 68 | logs: String::new(), 69 | status: match status { 70 | Health::Unknown | Health::Starting => Status::Starting, 71 | Health::SetupFailed => Status::Failed, 72 | _ => Status::Succeeded, 73 | }, 74 | started_at: Utc::now().to_rfc3339(), 75 | completed_at: Utc::now().to_rfc3339(), 76 | }, 77 | }) 78 | } 79 | 80 | #[allow(clippy::unused_async)] 81 | pub async fn shutdown(Extension(shutdown): Extension) -> Json { 82 | shutdown.start(); 83 | 84 | Json(String::new()) 85 | } 86 | -------------------------------------------------------------------------------- /lib/src/runner.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use atomic_enum::atomic_enum; 3 | use cog_core::{Cog, CogResponse}; 4 | use jsonschema::JSONSchema; 5 | use schemars::{schema_for, JsonSchema}; 6 | use serde_json::Value; 7 | use std::{ 8 | env, 9 | panic::{catch_unwind, AssertUnwindSafe}, 10 | sync::{atomic::Ordering, Arc, Mutex}, 11 | time::{Duration, Instant}, 12 | }; 13 | use tokio::sync::{mpsc, oneshot}; 14 | use tracing::{trace_span, Instrument}; 15 | 16 | use crate::{errors::ValidationErrorSet, shutdown::Shutdown}; 17 | 18 | #[derive(Debug, thiserror::Error)] 19 | pub enum Error { 20 | #[error("Runner is busy")] 21 | Busy, 22 | 23 | #[error("Prediction was canceled")] 24 | Canceled, 25 | 26 | #[error("The model panicked.")] 27 | Panic, 28 | 29 | #[error("Failed to validate input.")] 30 | Validation(ValidationErrorSet), 31 | 32 | #[error("Failed to run prediction: {0}")] 33 | Prediction(#[from] anyhow::Error), 34 | } 35 | 36 | #[atomic_enum] 37 | #[derive(serde::Serialize, JsonSchema)] 38 | #[serde(rename_all = "SCREAMING_SNAKE_CASE")] 39 | pub enum Health { 40 | Unknown, 41 | Starting, 42 | Ready, 43 | Busy, 44 | SetupFailed, 45 | } 46 | 47 | pub static RUNNER_HEALTH: AtomicHealth = AtomicHealth::new(Health::Unknown); 48 | 49 | type ResponseSender = oneshot::Sender>; 50 | 51 | #[derive(Clone)] 52 | pub struct Runner { 53 | schema: Arc, 54 | sender: mpsc::Sender<(ResponseSender, cog_core::http::Request)>, 55 | } 56 | 57 | impl Runner { 58 | pub fn new(shutdown: Shutdown, cancel: flume::Receiver<()>) -> Self { 59 | RUNNER_HEALTH.swap(Health::Starting, Ordering::SeqCst); 60 | 61 | let (sender, mut rx) = mpsc::channel::<(ResponseSender, cog_core::http::Request)>(1); 62 | 63 | let handle_shutdown = shutdown.clone(); 64 | let handle = tokio::spawn(async move { 65 | tracing::info!("Running setup()..."); 66 | let cog = tokio::select! { 67 | () = tokio::time::sleep(Duration::from_secs(5 * 60)) => { 68 | tracing::error!("Failed run setup(): Timed out"); 69 | RUNNER_HEALTH.swap(Health::SetupFailed, Ordering::SeqCst); 70 | handle_shutdown.start(); 71 | return; 72 | } 73 | cog = T::setup().instrument(trace_span!("cog_setup")) => { 74 | match cog { 75 | Ok(cog) => Arc::new(Mutex::new(cog)), 76 | Err(error) => { 77 | tracing::error!("Failed run setup(): {error}"); 78 | RUNNER_HEALTH.swap(Health::SetupFailed, Ordering::SeqCst); 79 | handle_shutdown.start(); 80 | return; 81 | } 82 | } 83 | } 84 | }; 85 | 86 | tracing::debug!("setup() finished. Cog is ready to accept predictions."); 87 | RUNNER_HEALTH.swap(Health::Ready, Ordering::SeqCst); 88 | if env::var("KUBERNETES_SERVICE_HOST").is_ok() { 89 | if let Err(err) = tokio::fs::create_dir_all("/var/run/cog").await { 90 | tracing::error!("Failed to create cog runtime state directory: {err}"); 91 | RUNNER_HEALTH.swap(Health::SetupFailed, Ordering::SeqCst); 92 | handle_shutdown.start(); 93 | return; 94 | } 95 | 96 | if let Err(error) = tokio::fs::File::create("/var/run/cog/ready").await { 97 | tracing::error!("Failed to signal cog is ready: {error}"); 98 | RUNNER_HEALTH.swap(Health::SetupFailed, Ordering::SeqCst); 99 | handle_shutdown.start(); 100 | return; 101 | } 102 | } 103 | 104 | // Cog is not Sync, so we wrap it with a Mutex and this function to run it from an async context (and thus make it cancellable). 105 | let run_prediction_async = |input| { 106 | async { 107 | let cog = cog.lock().unwrap(); 108 | 109 | catch_unwind(AssertUnwindSafe(|| cog.predict(input))) 110 | } 111 | .instrument(trace_span!("cog_predict")) 112 | }; 113 | 114 | while let Some((tx, req)) = rx.recv().await { 115 | tracing::debug!("Processing prediction: {req:?}"); 116 | RUNNER_HEALTH.swap(Health::Busy, Ordering::SeqCst); 117 | 118 | // We need spawn_blocking here to (sneakily) allow blocking code in serde Deserialize impls (used in `Path`, for example). 119 | let input = req.input.clone(); 120 | let input = 121 | tokio::task::spawn_blocking(move || serde_json::from_value(input).unwrap()) 122 | .await 123 | .unwrap(); 124 | 125 | let start = Instant::now(); 126 | tokio::select! { 127 | _ = cancel.recv_async() => { 128 | let _ = tx.send(Err(Error::Canceled)); 129 | tracing::debug!("Prediction canceled"); 130 | }, 131 | response = run_prediction_async(input) => { 132 | tracing::debug!("Prediction complete: {response:?}"); 133 | let _ = tx.send(match response { 134 | Err(_) => Err(Error::Panic), 135 | Ok(Err(error)) => Err(Error::Prediction(error)), 136 | Ok(Ok(response)) => match response.into_response(req).await { 137 | Err(error) => Err(Error::Prediction(error)), 138 | Ok(response) => Ok((response, start.elapsed())), 139 | }, 140 | }); 141 | } 142 | } 143 | 144 | RUNNER_HEALTH.swap(Health::Ready, Ordering::SeqCst); 145 | } 146 | }); 147 | 148 | tokio::spawn(async move { 149 | shutdown.handle().await; 150 | tracing::debug!("Shutting down runner..."); 151 | handle.abort(); 152 | }); 153 | 154 | let schema = jsonschema::JSONSchema::compile( 155 | &serde_json::to_value(schema_for!(T::Request)).unwrap(), 156 | ) 157 | .unwrap(); 158 | 159 | Self { 160 | sender, 161 | schema: Arc::new(schema), 162 | } 163 | } 164 | 165 | pub fn validate(&self, input: &Value) -> Result<(), ValidationErrorSet> { 166 | self.schema.validate(input)?; 167 | 168 | Ok(()) 169 | } 170 | 171 | pub async fn run(&self, req: cog_core::http::Request) -> Result<(Value, Duration), Error> { 172 | if !matches!(RUNNER_HEALTH.load(Ordering::SeqCst), Health::Ready) { 173 | tracing::debug!("Failed to run prediction: runner is busy"); 174 | return Err(Error::Busy); 175 | } 176 | 177 | self.validate(&req.input).map_err(Error::Validation)?; 178 | RUNNER_HEALTH.swap(Health::Busy, Ordering::SeqCst); 179 | 180 | let (tx, rx) = oneshot::channel(); 181 | 182 | tracing::debug!("Sending prediction to runner: {req:?}"); 183 | let _ = self.sender.send((tx, req)).await; 184 | tracing::debug!("Waiting for prediction response..."); 185 | let result = rx.await.unwrap(); 186 | tracing::debug!("Prediction response received: {result:?}"); 187 | 188 | RUNNER_HEALTH.swap(Health::Ready, Ordering::SeqCst); 189 | 190 | result 191 | } 192 | } 193 | -------------------------------------------------------------------------------- /lib/src/server.rs: -------------------------------------------------------------------------------- 1 | use std::{env, net::SocketAddr}; 2 | 3 | use aide::openapi::{self, OpenApi}; 4 | use anyhow::Result; 5 | use axum::{http::Method, Extension, Server}; 6 | use indexmap::indexmap; 7 | use schemars::{ 8 | gen::{SchemaGenerator, SchemaSettings}, 9 | schema::SchemaObject as Schema, 10 | }; 11 | 12 | use crate::{ 13 | helpers::openapi::{replace_request_schema, replace_response_schema, schema_with_properties}, 14 | prediction::Prediction, 15 | routes, 16 | shutdown::Shutdown, 17 | Cli, Cog, 18 | }; 19 | 20 | #[allow(clippy::redundant_pub_crate)] 21 | pub(crate) async fn start(args: Cli) -> Result<()> { 22 | if let Some(url) = args.upload_url { 23 | env::set_var("UPLOAD_URL", url.to_string()); 24 | } 25 | 26 | let mut openapi = generate_schema::(); 27 | let router = routes::handler().finish_api(&mut openapi); 28 | tweak_generated_schema(&mut openapi); 29 | 30 | let shutdown = Shutdown::new(args.await_explicit_shutdown.unwrap_or_default())?; 31 | if args.dump_schema_and_exit { 32 | println!("{}", serde_json::to_string(&openapi).unwrap()); 33 | shutdown.start(); 34 | return Ok(()); 35 | } 36 | 37 | let prediction = Prediction::setup::(shutdown.clone()); 38 | 39 | let router = router 40 | .layer(Extension(openapi)) 41 | .layer(shutdown.extension()) 42 | .layer(prediction.extension()); 43 | 44 | let addr = SocketAddr::from(( 45 | [0, 0, 0, 0], 46 | env::var("PORT").map_or(Ok(5000), |p| p.parse())?, 47 | )); 48 | 49 | tracing::info!("Starting server on {addr}..."); 50 | Server::bind(&addr) 51 | .serve(router.into_make_service()) 52 | .with_graceful_shutdown(shutdown.handle()) 53 | .await?; 54 | 55 | Ok(()) 56 | } 57 | 58 | fn generate_schema() -> OpenApi { 59 | let mut generator = SchemaGenerator::new(SchemaSettings::openapi3().with(|settings| { 60 | settings.inline_subschemas = true; 61 | })); 62 | 63 | OpenApi { 64 | info: openapi::Info { 65 | title: "Cog".to_string(), 66 | version: "0.1.0".to_string(), 67 | ..openapi::Info::default() 68 | }, 69 | components: Some(openapi::Components { 70 | schemas: indexmap! { 71 | "Input".to_string() => openapi::SchemaObject { 72 | example: None, 73 | external_docs: None, 74 | json_schema: schema_with_properties::(&mut generator, |name, schema, i| { 75 | schema.metadata().title = Some(titlecase::titlecase(&name)); 76 | schema.extensions.insert("x-order".to_string(), i.into()); 77 | }) 78 | }, 79 | "PredictionRequest".to_string() => openapi::SchemaObject { 80 | example: None, 81 | external_docs: None, 82 | json_schema: schema_with_properties::(&mut generator, |name, schema, _| { 83 | if name == "input" { 84 | schema.reference = Some("#/components/schemas/Input".to_string()); 85 | } 86 | }) 87 | }, 88 | "Output".to_string() => openapi::SchemaObject { 89 | example: None, 90 | external_docs: None, 91 | json_schema: generator.subschema_for::() 92 | }, 93 | "PredictionResponse".to_string() => openapi::SchemaObject { 94 | example: None, 95 | external_docs: None, 96 | json_schema: schema_with_properties::(&mut generator, |name, schema, _| { 97 | if name == "input" { 98 | schema.reference = Some("#/components/schemas/Input".to_string()); 99 | } 100 | 101 | if name == "output" { 102 | schema.reference = Some("#/components/schemas/Output".to_string()); 103 | } 104 | }) 105 | }, 106 | }, 107 | ..openapi::Components::default() 108 | }), 109 | ..OpenApi::default() 110 | } 111 | } 112 | 113 | fn tweak_generated_schema(openapi: &mut OpenApi) { 114 | replace_request_schema( 115 | openapi, 116 | "/predictions", 117 | (Method::POST, "application/json"), 118 | Schema::new_ref("#/components/schemas/PredictionRequest".to_string()), 119 | ) 120 | .unwrap(); 121 | 122 | replace_request_schema( 123 | openapi, 124 | "/predictions/{prediction_id}", 125 | (Method::PUT, "application/json"), 126 | Schema::new_ref("#/components/schemas/PredictionRequest".to_string()), 127 | ) 128 | .unwrap(); 129 | 130 | replace_response_schema( 131 | openapi, 132 | "/predictions", 133 | ( 134 | Method::POST, 135 | openapi::StatusCode::Code(200), 136 | "application/json".to_string(), 137 | ), 138 | Schema::new_ref("#/components/schemas/PredictionResponse".to_string()), 139 | ) 140 | .unwrap(); 141 | 142 | replace_response_schema( 143 | openapi, 144 | "/predictions/{prediction_id}", 145 | ( 146 | Method::PUT, 147 | openapi::StatusCode::Code(200), 148 | "application/json".to_string(), 149 | ), 150 | Schema::new_ref("#/components/schemas/PredictionResponse".to_string()), 151 | ) 152 | .unwrap(); 153 | } 154 | -------------------------------------------------------------------------------- /lib/src/shutdown.rs: -------------------------------------------------------------------------------- 1 | use axum::Extension; 2 | use std::{ 3 | error::Error, 4 | fmt, 5 | fmt::Display, 6 | future::Future, 7 | sync::atomic::{AtomicBool, Ordering}, 8 | }; 9 | use tokio::{signal, sync::broadcast}; 10 | 11 | #[derive(Debug, PartialEq, Eq)] 12 | pub struct AlreadyCreatedError; 13 | 14 | impl Error for AlreadyCreatedError {} 15 | 16 | impl Display for AlreadyCreatedError { 17 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 18 | f.write_str("shutdown handler already created") 19 | } 20 | } 21 | 22 | static CREATED: AtomicBool = AtomicBool::new(false); 23 | 24 | #[derive(Debug, Clone)] 25 | pub struct Shutdown { 26 | pub sender: broadcast::Sender<()>, 27 | } 28 | 29 | #[derive(Debug, Clone)] 30 | pub struct Agent { 31 | sender: broadcast::Sender<()>, 32 | } 33 | 34 | impl Agent { 35 | pub fn start(&self) { 36 | self.sender.send(()).ok(); 37 | } 38 | } 39 | 40 | impl Shutdown { 41 | pub fn new(await_explicit_shutdown: bool) -> Result { 42 | if (CREATED).swap(true, Ordering::SeqCst) { 43 | tracing::error!("shutdown handler called twice"); 44 | return Err(AlreadyCreatedError); 45 | } 46 | 47 | let (tx, _) = broadcast::channel(1); 48 | let handle = register_handlers(await_explicit_shutdown); 49 | 50 | let tx_for_handle = tx.clone(); 51 | tokio::spawn(async move { 52 | tracing::debug!("Registered shutdown handlers"); 53 | handle.await; 54 | tx_for_handle.send(()).ok(); 55 | }); 56 | 57 | Ok(Self { sender: tx }) 58 | } 59 | 60 | pub fn start(&self) { 61 | tracing::debug!("Manually requested shutdown."); 62 | self.sender.send(()).ok(); 63 | } 64 | 65 | pub fn handle(&self) -> impl Future + '_ { 66 | let mut rx = self.sender.subscribe(); 67 | 68 | async move { 69 | let rx = rx.recv(); 70 | 71 | rx.await.unwrap(); 72 | } 73 | } 74 | 75 | pub fn agent(&self) -> Agent { 76 | Agent { 77 | sender: self.sender.clone(), 78 | } 79 | } 80 | 81 | pub fn extension(&self) -> Extension { 82 | Extension(self.agent()) 83 | } 84 | } 85 | 86 | fn register_handlers(await_explicit_shutdown: bool) -> impl Future { 87 | let ctrl_c = async { 88 | signal::ctrl_c() 89 | .await 90 | .expect("failed to install Ctrl+C handler"); 91 | }; 92 | 93 | #[cfg(unix)] 94 | let terminate = async { 95 | signal::unix::signal(signal::unix::SignalKind::terminate()) 96 | .expect("failed to install signal handler") 97 | .recv() 98 | .await; 99 | }; 100 | 101 | #[cfg(not(unix))] 102 | let terminate = std::future::pending::<()>(); 103 | 104 | async move { 105 | if await_explicit_shutdown { 106 | return ctrl_c.await; 107 | } 108 | 109 | tokio::select! { 110 | () = ctrl_c => { 111 | tracing::info!("Received Ctrl+C signal"); 112 | }, 113 | () = terminate => { 114 | tracing::info!("Received terminate signal"); 115 | }, 116 | } 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /lib/src/spec.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use core::fmt::Debug; 3 | use mime_guess::Mime; 4 | use schemars::{gen::SchemaGenerator, schema::Schema, JsonSchema}; 5 | use serde::Serialize; 6 | use std::{ 7 | env::{self, temp_dir}, 8 | fs::File, 9 | path::PathBuf, 10 | str::FromStr, 11 | }; 12 | use url::Url; 13 | use uuid::Uuid; 14 | 15 | use crate::helpers::{base64_decode, base64_encode, url_join}; 16 | 17 | #[derive(Debug)] 18 | pub struct Path(PathBuf); 19 | 20 | impl Path { 21 | /// Create a new path from a url 22 | /// 23 | /// # Errors 24 | /// 25 | /// Returns an error if the url cannot be downloaded or a temporary file cannot be created. 26 | pub(crate) fn new(url: &Url) -> Result { 27 | if url.scheme() == "data" { 28 | return Self::from_dataurl(url); 29 | } 30 | 31 | tracing::debug!("Downloading file from {url}"); 32 | let file_path = temp_dir().join(url.path().split('/').last().unwrap_or_else(|| url.path())); 33 | let request = reqwest::blocking::get(url.as_str())?.bytes()?; 34 | 35 | std::io::copy(&mut request.as_ref(), &mut File::create(&file_path)?)?; 36 | tracing::debug!("Downloaded file to {}", file_path.display()); 37 | 38 | Ok(Self(file_path)) 39 | } 40 | 41 | /// Create a new path from a data url 42 | /// 43 | /// # Errors 44 | /// 45 | /// Returns an error if the url cannot be decoded or a temporary file cannot be created. 46 | pub(crate) fn from_dataurl(url: &Url) -> Result { 47 | let data = url.path().split(',').last().unwrap_or_else(|| url.path()); 48 | 49 | let file_bytes = base64_decode(data)?; 50 | let mime_type = Mime::from_str(tree_magic_mini::from_u8(&file_bytes)) 51 | .unwrap_or(mime_guess::mime::APPLICATION_OCTET_STREAM); 52 | let file_ext = mime_guess::get_mime_extensions(&mime_type) 53 | .and_then(<[&str]>::last) 54 | .map_or_else(String::new, |e| format!(".{e}")); 55 | 56 | let file_path = temp_dir().join(format!("{}{file_ext}", Uuid::new_v4())); 57 | 58 | std::fs::write(&file_path, file_bytes)?; 59 | Ok(Self(file_path)) 60 | } 61 | 62 | /// PUT the file to the given endpoint and return the url 63 | /// 64 | /// # Errors 65 | /// 66 | /// Returns an error if the file cannot be read or the upload fails. 67 | /// 68 | /// # Panics 69 | /// 70 | /// Panics if the file name is not valid unicode. 71 | pub(crate) fn upload_put(&self, upload_url: &Url) -> Result { 72 | let url = url_join(upload_url, self.0.file_name().unwrap().to_str().unwrap()); 73 | tracing::debug!("Uploading file to {url}"); 74 | 75 | let file_bytes = std::fs::read(&self.0)?; 76 | let mime_type = tree_magic_mini::from_u8(&file_bytes); 77 | 78 | let response = reqwest::blocking::Client::new() 79 | .put(url.clone()) 80 | .header("Content-Type", mime_type) 81 | .body(file_bytes) 82 | .send()?; 83 | 84 | if !response.status().is_success() { 85 | anyhow::bail!( 86 | "Failed to upload file to {url}: got {}. {}", 87 | response.status(), 88 | response.text().unwrap_or_default() 89 | ); 90 | } 91 | 92 | let mut url = response.url().clone(); 93 | url.set_query(None); 94 | 95 | tracing::debug!("Uploaded file to {url}"); 96 | Ok(url.to_string()) 97 | } 98 | 99 | /// Convert the file to a data url 100 | /// 101 | /// # Errors 102 | /// 103 | /// Returns an error if the file cannot be read. 104 | pub(crate) fn to_dataurl(&self) -> Result { 105 | let file_bytes = std::fs::read(&self.0)?; 106 | let mime_type = tree_magic_mini::from_u8(&file_bytes); 107 | 108 | Ok(format!( 109 | "data:{mime_type};base64,{base64}", 110 | base64 = base64_encode(&file_bytes) 111 | )) 112 | } 113 | } 114 | 115 | impl AsRef for Path { 116 | fn as_ref(&self) -> &std::path::Path { 117 | self.0.as_ref() 118 | } 119 | } 120 | 121 | impl JsonSchema for Path { 122 | fn schema_name() -> String { 123 | "Path".to_string() 124 | } 125 | 126 | fn json_schema(gen: &mut SchemaGenerator) -> Schema { 127 | Url::json_schema(gen) 128 | } 129 | } 130 | 131 | impl Drop for Path { 132 | fn drop(&mut self) { 133 | tracing::debug!("Removing temporary file at path {:?}", self.0); 134 | 135 | std::fs::remove_file(&self.0).unwrap(); 136 | } 137 | } 138 | 139 | impl<'de> serde::Deserialize<'de> for Path { 140 | fn deserialize(deserializer: D) -> std::result::Result 141 | where 142 | D: serde::Deserializer<'de>, 143 | { 144 | let url = String::deserialize(deserializer)?; 145 | 146 | Self::new(&Url::parse(&url).map_err(serde::de::Error::custom)?) 147 | .map_err(serde::de::Error::custom) 148 | } 149 | } 150 | 151 | impl Serialize for Path { 152 | fn serialize(&self, serializer: S) -> std::result::Result 153 | where 154 | S: serde::Serializer, 155 | { 156 | let url = env::var("UPLOAD_URL") 157 | .map(|url| url.parse().ok()) 158 | .ok() 159 | .flatten() 160 | .map_or_else( 161 | || self.to_dataurl(), 162 | |upload_url| self.upload_put(&upload_url), 163 | ); 164 | 165 | serializer.serialize_str(&url.map_err(serde::ser::Error::custom)?) 166 | } 167 | } 168 | 169 | impl From for Path { 170 | fn from(path: PathBuf) -> Self { 171 | Self(path) 172 | } 173 | } 174 | 175 | #[cfg(test)] 176 | mod tests { 177 | use super::*; 178 | use serde_json::json; 179 | 180 | #[derive(Debug, serde::Deserialize)] 181 | struct StructWithPath { 182 | file: Path, 183 | } 184 | 185 | #[test] 186 | fn test_path_deserialize() { 187 | let r#struct: StructWithPath = serde_json::from_value(json!({ 188 | "file": "https://raw.githubusercontent.com/m1guelpf/cog-rust/main/README.md" 189 | })) 190 | .unwrap(); 191 | 192 | let path = r#struct.file; 193 | let underlying_path = path.0.clone(); 194 | 195 | assert!( 196 | underlying_path.exists(), 197 | "File does not exist at path {:?}", 198 | path.0 199 | ); 200 | assert!( 201 | underlying_path.metadata().unwrap().len() > 0, 202 | "File is empty" 203 | ); 204 | 205 | drop(path); 206 | 207 | assert!( 208 | !underlying_path.exists(), 209 | "File still exists at path {underlying_path:?}", 210 | ); 211 | } 212 | 213 | #[test] 214 | fn test_dataurl_serialize() { 215 | let r#struct: StructWithPath = serde_json::from_value(json!({ 216 | "file": "https://upload.wikimedia.org/wikipedia/commons/thumb/1/1b/Square_200x200.png/120px-Square_200x200.png" 217 | })) 218 | .unwrap(); 219 | 220 | let path = r#struct.file; 221 | let dataurl = path.to_dataurl().unwrap(); 222 | 223 | assert!(dataurl.starts_with("data:image/png;base64,")); 224 | } 225 | } 226 | -------------------------------------------------------------------------------- /lib/src/webhooks.rs: -------------------------------------------------------------------------------- 1 | use std::env; 2 | 3 | use anyhow::Result; 4 | use axum::http::{HeaderMap, HeaderValue}; 5 | use cog_core::http::WebhookEvent; 6 | use reqwest::Client; 7 | use url::Url; 8 | 9 | use crate::prediction::{Prediction, ResponseHelpers}; 10 | 11 | pub struct WebhookSender { 12 | client: Client, 13 | } 14 | 15 | impl WebhookSender { 16 | pub fn new() -> Result { 17 | let mut headers = HeaderMap::new(); 18 | let client = Client::builder(); 19 | 20 | if let Ok(token) = env::var("WEBHOOK_AUTH_TOKEN") { 21 | let mut authorization = HeaderValue::from_str(&format!("Bearer {token}"))?; 22 | authorization.set_sensitive(true); 23 | headers.insert("Authorization", authorization); 24 | } 25 | 26 | Ok(Self { 27 | client: client 28 | .user_agent(format!("cog-worker/{}", env!("CARGO_PKG_VERSION"))) 29 | .default_headers(headers) 30 | .build()?, 31 | }) 32 | } 33 | 34 | pub async fn starting(&self, prediction: &Prediction) -> Result<()> { 35 | let request = prediction.request.clone().unwrap(); 36 | if !Self::should_send(&request, WebhookEvent::Start) { 37 | return Ok(()); 38 | } 39 | 40 | self.send( 41 | request.webhook.clone().unwrap(), 42 | cog_core::http::Response::starting(prediction.id.clone(), request), 43 | ) 44 | .await?; 45 | 46 | Ok(()) 47 | } 48 | 49 | pub async fn finished( 50 | &self, 51 | prediction: &Prediction, 52 | response: cog_core::http::Response, 53 | ) -> Result<()> { 54 | let request = prediction.request.clone().unwrap(); 55 | if !Self::should_send(&request, WebhookEvent::Completed) { 56 | return Ok(()); 57 | } 58 | 59 | self.send(request.webhook.clone().unwrap(), response) 60 | .await?; 61 | 62 | Ok(()) 63 | } 64 | 65 | fn should_send(req: &cog_core::http::Request, event: WebhookEvent) -> bool { 66 | req.webhook.is_some() 67 | && req 68 | .webhook_event_filters 69 | .as_ref() 70 | .map_or(true, |filters| filters.contains(&event)) 71 | } 72 | 73 | async fn send( 74 | &self, 75 | url: Url, 76 | res: cog_core::http::Response, 77 | ) -> Result { 78 | tracing::debug!("Sending webhook to {url}"); 79 | tracing::trace!("{res:?}"); 80 | 81 | self.client.post(url).json(&res).send().await 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | tab_spaces = 4 2 | hard_tabs = true 3 | edition = "2021" 4 | use_try_shorthand = true 5 | imports_granularity = "crate" 6 | use_field_init_shorthand = true 7 | condense_wildcard_suffixes = true 8 | match_block_trailing_comma = true 9 | --------------------------------------------------------------------------------