feat(common): replace need_scheduling bool with SchedulingPolicy enum in Download

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dragonfly-api"
-version = "2.3.4"
+version = "2.3.5"
 authors = ["Gaius <gaius.qi@gmail.com>"]
 edition = "2021"
 license = "Apache-2.0"

pkg/apis/common/v2/common.proto

@@ -119,6 +119,19 @@ enum Priority {
   LEVEL6 = 6;
 }
 
+// SchedulingPolicy represents how the download interacts with the scheduler.
+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.
+  AUTO = 0;
+
+  // 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.
+  ALWAYS = 1;
+}
+
 // Peer metadata.
 message Peer {
   // Peer id.
@@ -727,11 +740,9 @@ message Download {
   bool metadata_only = 33;
   // OpenCSG protocol information.
   optional OpenCSG open_csg = 34;
-  // Need scheduling is the flag to indicate whether the download needs to be scheduled
-  // by 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.
-  bool need_scheduling = 35;
+  // Scheduling policy represents how the download interacts with the scheduler,
+  // default is AUTO.
+  SchedulingPolicy scheduling_policy = 35 [(validate.rules).enum.defined_only = true];
 
   reserved 21;
   reserved "load_to_cache";

proto/common.proto

@@ -113,6 +113,19 @@ enum Priority {
   LEVEL6 = 6;
 }
 
+// SchedulingPolicy represents how the download interacts with the scheduler.
+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.
+  AUTO = 0;
+
+  // 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.
+  ALWAYS = 1;
+}
+
 // Peer metadata.
 message Peer {
   // Peer id.
@@ -669,11 +682,9 @@ message Download {
   bool metadata_only = 33;
   // OpenCSG protocol information.
   optional OpenCSG open_csg = 34;
-  // Need scheduling is the flag to indicate whether the download needs to be scheduled
-  // by 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.
-  bool need_scheduling = 35;
+  // Scheduling policy represents how the download interacts with the scheduler,
+  // default is AUTO.
+  SchedulingPolicy scheduling_policy = 35;
 
   reserved 21;
   reserved "load_to_cache";

python/dragonfly_api/common_pb2.pyi

@@ -40,6 +40,11 @@ class Priority(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     LEVEL4: _ClassVar[Priority]
     LEVEL5: _ClassVar[Priority]
     LEVEL6: _ClassVar[Priority]
+
+class SchedulingPolicy(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    AUTO: _ClassVar[SchedulingPolicy]
+    ALWAYS: _ClassVar[SchedulingPolicy]
 NORMAL: SizeScope
 SMALL: SizeScope
 TINY: SizeScope
@@ -58,6 +63,8 @@ LEVEL3: Priority
 LEVEL4: Priority
 LEVEL5: Priority
 LEVEL6: Priority
+AUTO: SchedulingPolicy
+ALWAYS: SchedulingPolicy
 
 class Peer(_message.Message):
     __slots__ = ("id", "range", "priority", "pieces", "cost", "state", "task", "host", "need_back_to_source", "created_at", "updated_at", "concurrent_piece_count")
@@ -494,7 +501,7 @@ class Build(_message.Message):
     def __init__(self, git_version: _Optional[str] = ..., git_commit: _Optional[str] = ..., go_version: _Optional[str] = ..., rust_version: _Optional[str] = ..., platform: _Optional[str] = ...) -> None: ...
 
 class Download(_message.Message):
-    __slots__ = ("url", "digest", "range", "type", "tag", "application", "priority", "filtered_query_params", "request_header", "piece_length", "output_path", "timeout", "disable_back_to_source", "need_back_to_source", "certificate_chain", "prefetch", "object_storage", "hdfs", "is_prefetch", "need_piece_content", "force_hard_link", "content_for_calculating_task_id", "remote_ip", "concurrent_piece_count", "overwrite", "actual_piece_length", "actual_content_length", "actual_piece_count", "enable_task_id_based_blob_digest", "hugging_face", "model_scope", "metadata_only", "open_csg", "need_scheduling")
+    __slots__ = ("url", "digest", "range", "type", "tag", "application", "priority", "filtered_query_params", "request_header", "piece_length", "output_path", "timeout", "disable_back_to_source", "need_back_to_source", "certificate_chain", "prefetch", "object_storage", "hdfs", "is_prefetch", "need_piece_content", "force_hard_link", "content_for_calculating_task_id", "remote_ip", "concurrent_piece_count", "overwrite", "actual_piece_length", "actual_content_length", "actual_piece_count", "enable_task_id_based_blob_digest", "hugging_face", "model_scope", "metadata_only", "open_csg", "scheduling_policy")
     class RequestHeaderEntry(_message.Message):
         __slots__ = ("key", "value")
         KEY_FIELD_NUMBER: _ClassVar[int]
@@ -535,7 +542,7 @@ class Download(_message.Message):
     MODEL_SCOPE_FIELD_NUMBER: _ClassVar[int]
     METADATA_ONLY_FIELD_NUMBER: _ClassVar[int]
     OPEN_CSG_FIELD_NUMBER: _ClassVar[int]
-    NEED_SCHEDULING_FIELD_NUMBER: _ClassVar[int]
+    SCHEDULING_POLICY_FIELD_NUMBER: _ClassVar[int]
     url: str
     digest: str
     range: Range
@@ -569,8 +576,8 @@ class Download(_message.Message):
     model_scope: ModelScope
     metadata_only: bool
     open_csg: OpenCSG
-    need_scheduling: bool
-    def __init__(self, url: _Optional[str] = ..., digest: _Optional[str] = ..., range: _Optional[_Union[Range, _Mapping]] = ..., type: _Optional[_Union[TaskType, str]] = ..., tag: _Optional[str] = ..., application: _Optional[str] = ..., priority: _Optional[_Union[Priority, str]] = ..., filtered_query_params: _Optional[_Iterable[str]] = ..., request_header: _Optional[_Mapping[str, str]] = ..., piece_length: _Optional[int] = ..., output_path: _Optional[str] = ..., timeout: _Optional[_Union[datetime.timedelta, _duration_pb2.Duration, _Mapping]] = ..., disable_back_to_source: bool = ..., need_back_to_source: bool = ..., certificate_chain: _Optional[_Iterable[bytes]] = ..., prefetch: bool = ..., object_storage: _Optional[_Union[ObjectStorage, _Mapping]] = ..., hdfs: _Optional[_Union[HDFS, _Mapping]] = ..., is_prefetch: bool = ..., need_piece_content: bool = ..., force_hard_link: bool = ..., content_for_calculating_task_id: _Optional[str] = ..., remote_ip: _Optional[str] = ..., concurrent_piece_count: _Optional[int] = ..., overwrite: bool = ..., actual_piece_length: _Optional[int] = ..., actual_content_length: _Optional[int] = ..., actual_piece_count: _Optional[int] = ..., enable_task_id_based_blob_digest: bool = ..., hugging_face: _Optional[_Union[HuggingFace, _Mapping]] = ..., model_scope: _Optional[_Union[ModelScope, _Mapping]] = ..., metadata_only: bool = ..., open_csg: _Optional[_Union[OpenCSG, _Mapping]] = ..., need_scheduling: bool = ...) -> None: ...
+    scheduling_policy: SchedulingPolicy
+    def __init__(self, url: _Optional[str] = ..., digest: _Optional[str] = ..., range: _Optional[_Union[Range, _Mapping]] = ..., type: _Optional[_Union[TaskType, str]] = ..., tag: _Optional[str] = ..., application: _Optional[str] = ..., priority: _Optional[_Union[Priority, str]] = ..., filtered_query_params: _Optional[_Iterable[str]] = ..., request_header: _Optional[_Mapping[str, str]] = ..., piece_length: _Optional[int] = ..., output_path: _Optional[str] = ..., timeout: _Optional[_Union[datetime.timedelta, _duration_pb2.Duration, _Mapping]] = ..., disable_back_to_source: bool = ..., need_back_to_source: bool = ..., certificate_chain: _Optional[_Iterable[bytes]] = ..., prefetch: bool = ..., object_storage: _Optional[_Union[ObjectStorage, _Mapping]] = ..., hdfs: _Optional[_Union[HDFS, _Mapping]] = ..., is_prefetch: bool = ..., need_piece_content: bool = ..., force_hard_link: bool = ..., content_for_calculating_task_id: _Optional[str] = ..., remote_ip: _Optional[str] = ..., concurrent_piece_count: _Optional[int] = ..., overwrite: bool = ..., actual_piece_length: _Optional[int] = ..., actual_content_length: _Optional[int] = ..., actual_piece_count: _Optional[int] = ..., enable_task_id_based_blob_digest: bool = ..., hugging_face: _Optional[_Union[HuggingFace, _Mapping]] = ..., model_scope: _Optional[_Union[ModelScope, _Mapping]] = ..., metadata_only: bool = ..., open_csg: _Optional[_Union[OpenCSG, _Mapping]] = ..., scheduling_policy: _Optional[_Union[SchedulingPolicy, str]] = ...) -> None: ...
 
 class ObjectStorage(_message.Message):
     __slots__ = ("region", "endpoint", "access_key_id", "access_key_secret", "session_token", "credential_path", "predefined_acl", "security_token", "insecure_skip_verify")

src/common.v2.rs

@@ -784,12 +784,10 @@ pub struct Download {
     /// OpenCSG protocol information.
     #[prost(message, optional, tag = "34")]
     pub open_csg: ::core::option::Option<OpenCsg>,
-    /// Need scheduling is the flag to indicate whether the download needs to be scheduled
-    /// by 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.
-    #[prost(bool, tag = "35")]
-    pub need_scheduling: bool,
+    /// Scheduling policy represents how the download interacts with the scheduler,
+    /// default is AUTO.
+    #[prost(enumeration = "SchedulingPolicy", tag = "35")]
+    pub scheduling_policy: i32,
 }
 /// Object Storage related information.
 #[derive(serde::Serialize, serde::Deserialize)]
@@ -1105,3 +1103,37 @@ impl Priority {
         }
     }
 }
+/// SchedulingPolicy represents how the download interacts with the scheduler.
+#[derive(serde::Serialize, serde::Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+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.
+    Auto = 0,
+    /// 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.
+    Always = 1,
+}
+impl SchedulingPolicy {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            Self::Auto => "AUTO",
+            Self::Always => "ALWAYS",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "AUTO" => Some(Self::Auto),
+            "ALWAYS" => Some(Self::Always),
+            _ => None,
+        }
+    }
+}