tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

feat: enable azure and oci remote file system

Browse Source
This commit is contained in:
NanoCode012
2025-04-04 19:04:24 +07:00
parent a945abbeae
commit df396ed303
1 changed files with 41 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
src/axolotl/utils/data/shared.py
View File

@@ -96,20 +96,17 @@ def load_dataset_w_config(
pass pass
ds_from_cloud = False ds_from_cloud = False
storage_options = {} storage_options: dict = {}
remote_file_system = None remote_file_system = None
if config_dataset.path.startswith("s3://"): if config_dataset.path.startswith("s3://"):
try: try:
import aiobotocore.session # type: ignore
import s3fs # type: ignore import s3fs # type: ignore
except ImportError as exc: except ImportError as exc:
raise ImportError( raise ImportError("s3:// paths require s3fs to be installed") from exc
"s3:// paths require aiobotocore and s3fs to be installed"
) from exc
# Takes credentials from ~/.aws/credentials for default profile # Reads env, credentials from ~/.aws/credentials, or IAM metadata provider
s3_session = aiobotocore.session.AioSession(profile="default") # https://s3fs.readthedocs.io/en/latest/index.html?highlight=storage_options#credentials
storage_options = {"session": s3_session} storage_options = {"anon": False}
remote_file_system = s3fs.S3FileSystem(**storage_options) remote_file_system = s3fs.S3FileSystem(**storage_options)
elif config_dataset.path.startswith("gs://") or config_dataset.path.startswith( elif config_dataset.path.startswith("gs://") or config_dataset.path.startswith(
"gcs://" "gcs://"
@@ -125,28 +122,44 @@ def load_dataset_w_config(
# https://gcsfs.readthedocs.io/en/latest/#credentials # https://gcsfs.readthedocs.io/en/latest/#credentials
storage_options = {"token": None} storage_options = {"token": None}
remote_file_system = gcsfs.GCSFileSystem(**storage_options) remote_file_system = gcsfs.GCSFileSystem(**storage_options)
# TODO: Figure out how to get auth creds passed elif (
# elif config_dataset.path.startswith("adl://") or config_dataset.path.startswith("abfs://"): config_dataset.path.startswith("adl://")
# try: or config_dataset.path.startswith("abfs://")
# import adlfs or config_dataset.path.startswith("az://")
# except ImportError as exc: ):
# raise ImportError( try:
# "adl:// or abfs:// paths require adlfs to be installed" import adlfs
# ) from exc except ImportError as exc:
raise ImportError(
"adl:// or abfs:// paths require adlfs to be installed"
) from exc
# # Gen 1 # # Ensure you have the following environment variables set:
# storage_options = { # # Gen 1
# "tenant_id": TENANT_ID, # storage_options = {
# "client_id": CLIENT_ID, # "tenant_id": AZURE_STORAGE_TENANT_ID,
# "client_secret": CLIENT_SECRET, # "client_id": AZURE_STORAGE_CLIENT_ID,
# } # "client_secret": AZURE_STORAGE_CLIENT_SECRET,
# # Gen 2 # }
# storage_options = { # # Gen 2
# "account_name": ACCOUNT_NAME, # storage_options = {
# "account_key": ACCOUNT_KEY, # "account_name": AZURE_STORAGE_ACCOUNT_NAME,
# } # "account_key": AZURE_STORAGE_ACCOUNT_KEY,
# }
# Reads env
# https://github.com/fsspec/adlfs?tab=readme-ov-file#setting-credentials
storage_options = {"anon": False}
remote_file_system = adlfs.AzureBlobFileSystem(**storage_options)
elif config_dataset.path.startswith("oci://"):
try:
import ocifs
except ImportError as exc:
raise ImportError("oci:// paths require ocifs to be installed") from exc
# https://ocifs.readthedocs.io/en/latest/getting-connected.html#Using-Environment-Variables
remote_file_system = ocifs.OCIFileSystem(**storage_options)
# remote_file_system = adlfs.AzureBlobFileSystem(**storage_options)
try: try:
if remote_file_system and remote_file_system.exists(config_dataset.path): if remote_file_system and remote_file_system.exists(config_dataset.path):
ds_from_cloud = True ds_from_cloud = True