chore: release v1.5.3 with scheduling policy support

Cargo.toml

@@ -13,7 +13,7 @@ members = [
 ]
 
 [workspace.package]
-version = "1.5.2"
+version = "1.5.3"
 authors = ["The Dragonfly Developers"]
 edition = "2021"
 readme = "README.md"
@@ -33,15 +33,15 @@ chrono = { version = "0.4.45", features = ["clock", "serde"] }
 clap = { version = "4.6.1", features = ["derive", "env"] }
 crc32fast = "1.5.0"
 dashmap = "6.1.0"
-dragonfly-api = "=2.3.1"
-dragonfly-client = { path = "dragonfly-client", version = "1.5.2" }
-dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.5.2" }
-dragonfly-client-config = { path = "dragonfly-client-config", version = "1.5.2" }
-dragonfly-client-core = { path = "dragonfly-client-core", version = "1.5.2" }
-dragonfly-client-init = { path = "dragonfly-client-init", version = "1.5.2" }
-dragonfly-client-metric = { path = "dragonfly-client-metric", version = "1.5.2" }
-dragonfly-client-storage = { path = "dragonfly-client-storage", version = "1.5.2" }
-dragonfly-client-util = { path = "dragonfly-client-util", version = "1.5.2" }
+dragonfly-api = "=2.3.5"
+dragonfly-client = { path = "dragonfly-client", version = "1.5.3" }
+dragonfly-client-backend = { path = "dragonfly-client-backend", version = "1.5.3" }
+dragonfly-client-config = { path = "dragonfly-client-config", version = "1.5.3" }
+dragonfly-client-core = { path = "dragonfly-client-core", version = "1.5.3" }
+dragonfly-client-init = { path = "dragonfly-client-init", version = "1.5.3" }
+dragonfly-client-metric = { path = "dragonfly-client-metric", version = "1.5.3" }
+dragonfly-client-storage = { path = "dragonfly-client-storage", version = "1.5.3" }
+dragonfly-client-util = { path = "dragonfly-client-util", version = "1.5.3" }
 fastrand = "2.4.1"
 fs2 = "0.4.3"
 futures = "0.3.32"

dragonfly-client-config/Cargo.toml

@@ -11,6 +11,7 @@ edition.workspace = true
 build = "build.rs"
 
 [dependencies]
+dragonfly-api.workspace = true
 dragonfly-client-core.workspace = true
 dragonfly-client-util.workspace = true
 local-ip-address.workspace = true

dragonfly-client-config/src/dfdaemon.rs

@@ -1300,6 +1300,33 @@ impl ProxyServer {
     }
 }
 
+/// SchedulingPolicy represents how the download interacts with the scheduler.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum SchedulingPolicy {
+    /// Auto downloads through the scheduler unless the content length is smaller
+    /// than the minimum piece length, in which case it downloads from the source
+    /// directly, skipping the scheduler.
+    #[default]
+    Auto,
+
+    /// Always downloads through the scheduler even if the content length is smaller
+    /// than the minimum piece length, so that the peer announces the task to the
+    /// scheduler and other peers can discover it as a parent. It is useful for
+    /// sharing small artifacts, such as OCI image manifests.
+    Always,
+}
+
+/// Implement From<SchedulingPolicy> for the scheduling policy of the download.
+impl From<SchedulingPolicy> for dragonfly_api::common::v2::SchedulingPolicy {
+    fn from(policy: SchedulingPolicy) -> Self {
+        match policy {
+            SchedulingPolicy::Auto => dragonfly_api::common::v2::SchedulingPolicy::Auto,
+            SchedulingPolicy::Always => dragonfly_api::common::v2::SchedulingPolicy::Always,
+        }
+    }
+}
+
 /// The proxy rule configuration.
 #[derive(Debug, Clone, Validate, Deserialize)]
 #[serde(default, rename_all = "camelCase")]
@@ -1322,6 +1349,12 @@ pub struct Rule {
     /// Default value includes the filtered query params of s3, gcs, oss, obs, cos.
     #[serde(default = "default_proxy_rule_filtered_query_params")]
     pub filtered_query_params: Vec<String>,
+
+    /// Represents how the download interacts with the scheduler, default is auto.
+    /// Auto downloads small files from the source directly, skipping the scheduler.
+    /// Always downloads through the scheduler even for small files, so that the peer
+    /// announces the task to the scheduler and other peers can discover it as a parent.
+    pub scheduling_policy: SchedulingPolicy,
 }
 
 /// Implement Default for Rule.
@@ -1332,6 +1365,7 @@ impl Default for Rule {
             use_tls: false,
             redirect: None,
             filtered_query_params: default_proxy_rule_filtered_query_params(),
+            scheduling_policy: SchedulingPolicy::default(),
         }
     }
 }
@@ -2055,6 +2089,26 @@ mod tests {
         assert!(config.is_some());
     }
 
+    #[test]
+    fn deserialize_proxy_rule_correctly() {
+        let yaml = r#"
+regex: 'manifests/sha256:.*'
+schedulingPolicy: always
+"#;
+
+        let rule: Rule = serde_yaml::from_str(yaml).expect("Failed to deserialize");
+        assert!(rule.regex.is_match("https://example.com/v2/library/ubuntu/manifests/sha256:b2c366cce7e68013d5441c6326d5a3e1b12aeb5ed58564d0fd3fa089bc29cb6e"));
+        assert_eq!(rule.scheduling_policy, SchedulingPolicy::Always);
+
+        let yaml = r#"
+regex: 'blobs/sha256.*'
+"#;
+
+        let rule: Rule = serde_yaml::from_str(yaml).expect("Failed to deserialize");
+        assert_eq!(rule.scheduling_policy, SchedulingPolicy::Auto);
+        assert_eq!(Rule::default().scheduling_policy, SchedulingPolicy::Auto);
+    }
+
     #[test]
     fn deserialize_optional_fields_correctly() {
         let yaml = r#"

dragonfly-client-util/src/types/redacted.rs

@@ -83,6 +83,7 @@ impl fmt::Debug for RedactedDownload<'_> {
             timeout: _,
             disable_back_to_source: _,
             need_back_to_source: _,
+            scheduling_policy: _,
             certificate_chain: _,
             need_piece_content: _,
             force_hard_link: _,

dragonfly-client/src/bin/dfget/main.rs

@@ -17,7 +17,7 @@
 use bytesize::ByteSize;
 use clap::Parser;
 use dragonfly_api::common::v2::{
-    Download, Hdfs, HuggingFace, ModelScope, ObjectStorage, OpenCsg, TaskType,
+    Download, Hdfs, HuggingFace, ModelScope, ObjectStorage, OpenCsg, SchedulingPolicy, TaskType,
 };
 use dragonfly_api::dfdaemon::v2::{
     download_task_response, DownloadTaskRequest, ListTaskEntriesRequest,
@@ -264,6 +264,15 @@ struct Args {
     )]
     disable_back_to_source: bool,
 
+    #[arg(
+        long = "scheduling-policy",
+        default_value = "auto",
+        value_parser = ["auto", "always"],
+        env = "DFGET_SCHEDULING_POLICY",
+        help = "Specify how the download interacts with the scheduler, 'auto' downloads small files from the source directly, 'always' downloads through the scheduler even for small files, so the peer announces the task and other peers can discover it as a parent"
+    )]
+    scheduling_policy: String,
+
     #[arg(
         long,
         env = "DFGET_STORAGE_REGION",
@@ -964,6 +973,11 @@ async fn download(
                         .or_err(ErrorType::ParseError)?,
                 ),
                 need_back_to_source: false,
+                scheduling_policy: if args.scheduling_policy.eq_ignore_ascii_case("always") {
+                    SchedulingPolicy::Always as i32
+                } else {
+                    SchedulingPolicy::Auto as i32
+                },
                 disable_back_to_source: args.disable_back_to_source,
                 certificate_chain: Vec::new(),
                 prefetch: false,

dragonfly-client/src/proxy/header.rs

@@ -15,7 +15,7 @@
  */
 
 use bytesize::ByteSize;
-use dragonfly_api::common::v2::Priority;
+use dragonfly_api::common::v2::{Priority, SchedulingPolicy};
 use reqwest::header::HeaderMap;
 use std::{fmt, str::FromStr};
 use tracing::error;
@@ -88,6 +88,14 @@ pub const DRAGONFLY_CONTENT_FOR_CALCULATING_TASK_ID_HEADER: &str =
 pub const DRAGONFLY_ENABLE_TASK_ID_BASED_BLOB_DIGEST: &str =
     "X-Dragonfly-Enable-Task-ID-Based-Blob-Digest";
 
+/// The header key of scheduling policy. It represents how the download interacts
+/// with the scheduler. The value is case-insensitive, e.g. "auto" or "always".
+/// "auto" downloads small files from the source directly, skipping the scheduler.
+/// "always" downloads through the scheduler even if the content length is smaller
+/// than the minimum piece length, so that the peer announces the task to the
+/// scheduler and other peers can discover it as a parent.
+pub const DRAGONFLY_SCHEDULING_POLICY_HEADER: &str = "X-Dragonfly-Scheduling-Policy";
+
 /// The response header key to indicate whether the task download finished.
 /// When the task download is finished, the response will include this header with the value `"true"`,
 /// indicating that the download hit the local cache.
@@ -322,6 +330,26 @@ pub fn get_enable_task_id_based_blob_digest(header: &HeaderMap, default: bool) -
     }
 }
 
+/// Get X-Dragonfly-Scheduling-Policy header value to determine how the download
+/// interacts with the scheduler.
+pub fn get_scheduling_policy(header: &HeaderMap, default: SchedulingPolicy) -> SchedulingPolicy {
+    match header.get(DRAGONFLY_SCHEDULING_POLICY_HEADER) {
+        Some(value) => match value.to_str() {
+            Ok(value) if value.eq_ignore_ascii_case("auto") => SchedulingPolicy::Auto,
+            Ok(value) if value.eq_ignore_ascii_case("always") => SchedulingPolicy::Always,
+            Ok(value) => {
+                error!("invalid scheduling policy from header: {}", value);
+                default
+            }
+            Err(err) => {
+                error!("get scheduling policy from header failed: {}", err);
+                default
+            }
+        },
+        None => default,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -517,4 +545,47 @@ mod tests {
         assert!(get_enable_task_id_based_blob_digest(&empty_headers, true));
         assert!(!get_enable_task_id_based_blob_digest(&empty_headers, false));
     }
+
+    #[test]
+    fn test_get_scheduling_policy() {
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            DRAGONFLY_SCHEDULING_POLICY_HEADER,
+            HeaderValue::from_static("always"),
+        );
+        assert_eq!(
+            get_scheduling_policy(&headers, SchedulingPolicy::Auto),
+            SchedulingPolicy::Always
+        );
+
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            DRAGONFLY_SCHEDULING_POLICY_HEADER,
+            HeaderValue::from_static("AUTO"),
+        );
+        assert_eq!(
+            get_scheduling_policy(&headers, SchedulingPolicy::Always),
+            SchedulingPolicy::Auto
+        );
+
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            DRAGONFLY_SCHEDULING_POLICY_HEADER,
+            HeaderValue::from_static("invalid"),
+        );
+        assert_eq!(
+            get_scheduling_policy(&headers, SchedulingPolicy::Always),
+            SchedulingPolicy::Always
+        );
+
+        let empty_headers = HeaderMap::new();
+        assert_eq!(
+            get_scheduling_policy(&empty_headers, SchedulingPolicy::Always),
+            SchedulingPolicy::Always
+        );
+        assert_eq!(
+            get_scheduling_policy(&empty_headers, SchedulingPolicy::Auto),
+            SchedulingPolicy::Auto
+        );
+    }
 }

dragonfly-client/src/proxy/mod.rs

@@ -1304,6 +1304,8 @@ fn make_download_task_request(
             output_path: header::get_output_path(header),
             timeout: None,
             need_back_to_source: false,
+            scheduling_policy: header::get_scheduling_policy(header, rule.scheduling_policy.into())
+                as i32,
             disable_back_to_source: config.proxy.disable_back_to_source,
             certificate_chain: Vec::new(),
             prefetch: need_prefetch(&config, header),