feat(trainer): add s3 initializers, add ignore_patterns to hf initializers

Signed-off-by: rudeigerc <[email protected]>

Author: rudeigerc

kubeflow/trainer/types/types.py

@@ -277,12 +277,42 @@ class TrainJob:
 # TODO (andreyvelich): Discuss how to keep these configurations is sync with pkg.initializers.types
 @dataclass
 class HuggingFaceDatasetInitializer:
-    """Configuration for downloading datasets from HuggingFace Hub."""
+    """Configuration for downloading datasets from HuggingFace Hub.
+
+    Args:
+        storage_uri (`str`): The HuggingFace Hub model identifier in the format 'hf://username/repo_name'.
+        ignore_patterns (`Optional[list[str]]`): List of file patterns to ignore during download.
+        access_token (`Optional[str]`): HuggingFace Hub access token for private datasets.
+    """
 
     storage_uri: str
+    ignore_patterns: Optional[list[str]] = None
     access_token: Optional[str] = None
 
 
+@dataclass
+class S3DatasetInitializer:
+    """Configuration for downloading datasets from S3-compatible storage.
+
+    Args:
+        storage_uri (`str`): The S3 URI for the model in the format 's3://bucket-name/path/to/model'.
+        ignore_patterns (`Optional[list[str]]`): List of file patterns to ignore during download.
+        endpoint (`Optional[str]`): Custom S3 endpoint URL.
+        access_key_id (`Optional[str]`): Access key for authentication.
+        secret_access_key (`Optional[str]`): Secret key for authentication.
+        region (`Optional[str]`): Region used in instantiating the client.
+        role_arn (`Optional[str]`): The ARN of the role you want to assume.
+    """
+
+    storage_uri: str
+    ignore_patterns: Optional[list[str]] = None
+    endpoint: Optional[str] = None
+    access_key_id: Optional[str] = None
+    secret_access_key: Optional[str] = None
+    region: Optional[str] = None
+    role_arn: Optional[str] = None
+
+
 @dataclass
 class DataCacheInitializer:
     """Configuration for distributed data caching system for training workloads.
@@ -332,23 +362,60 @@ def __post_init__(self):
 # Configuration for the HuggingFace model initializer.
 @dataclass
 class HuggingFaceModelInitializer:
+    """Configuration for downloading models from HuggingFace Hub.
+
+    Args:
+        storage_uri (`str`): The HuggingFace Hub model identifier in the format 'hf://username/repo_name'.
+        ignore_patterns (`Optional[list[str]]`): List of file patterns to ignore during download.
+        access_token (`Optional[str]`): HuggingFace Hub access token.
+    """
+
     storage_uri: str
+    ignore_patterns: Optional[list[str]] = None
     access_token: Optional[str] = None
 
 
+@dataclass
+class S3ModelInitializer:
+    """Configuration for downloading models from S3-compatible storage.
+
+    Args:
+        storage_uri (`str`): The S3 URI for the model in the format 's3://bucket-name/path/to/model'.
+        ignore_patterns (`Optional[list[str]]`): List of file patterns to ignore during download.
+            Defaults to `['*.msgpack', '*.h5', '*.bin', '.pt', '.pth']`.
+        endpoint (`Optional[str]`): Custom S3 endpoint URL.
+        access_key_id (`Optional[str]`): Access key for authentication.
+        secret_access_key (`Optional[str]`): Secret key for authentication.
+        region (`Optional[str]`): Region used in instantiating the client.
+        role_arn (`Optional[str]`): The ARN of the role you want to assume.
+    """
+
+    storage_uri: str
+    ignore_patterns: Optional[list[str]] = field(
+        default_factory=lambda: ["*.msgpack", "*.h5", "*.bin", ".pt", ".pth"]
+    )
+    endpoint: Optional[str] = None
+    access_key_id: Optional[str] = None
+    secret_access_key: Optional[str] = None
+    region: Optional[str] = None
+    role_arn: Optional[str] = None
+
+
 @dataclass
 class Initializer:
     """Initializer defines configurations for dataset and pre-trained model initialization
 
     Args:
-        dataset (`Optional[Union[HuggingFaceDatasetInitializer, DataCacheInitializer]]`):
+        dataset (`Optional[Union[HuggingFaceDatasetInitializer, S3DatasetInitializer, DataCacheInitializer]]`):
             The configuration for one of the supported dataset initializers.
-        model (`Optional[HuggingFaceModelInitializer]`): The configuration for one of the
-            supported model initializers.
-    """
-
-    dataset: Optional[Union[HuggingFaceDatasetInitializer, DataCacheInitializer]] = None
-    model: Optional[HuggingFaceModelInitializer] = None
+        model (`Optional[Union[HuggingFaceModelInitializer, S3ModelInitializer]]`):
+            The configuration for one of the supported model initializers.
+    """  # noqa: E501
+
+    dataset: Optional[
+        Union[HuggingFaceDatasetInitializer, S3DatasetInitializer, DataCacheInitializer]
+    ] = None
+    model: Optional[Union[HuggingFaceModelInitializer, S3ModelInitializer]] = None
 
 
 # TODO (andreyvelich): Add train() and optimize() methods to this class.