mirror of
https://github.com/hanzo-ml/kserve.git
synced 2026-10-30 20:53:49 +00:00
Signed-off-by: Andres Llausas <[email protected]> Co-authored-by: Claude Sonnet 4.6 <[email protected]>
1269 lines
48 KiB
Python
1269 lines
48 KiB
Python
# Copyright 2023 The KServe Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import asyncio
|
|
import base64
|
|
import fnmatch
|
|
from functools import partial
|
|
import glob
|
|
import gzip
|
|
import json
|
|
import mimetypes
|
|
import multiprocessing
|
|
import os
|
|
import re
|
|
import shutil
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
from typing import List, Optional
|
|
import zipfile
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
from urllib.parse import urlparse
|
|
import requests
|
|
|
|
from kserve_storage.logging import logger
|
|
from kserve_storage.storage_errors import (
|
|
raise_storage_error,
|
|
check_http_response,
|
|
)
|
|
|
|
MODEL_MOUNT_DIRS = "/mnt/models"
|
|
|
|
_GCS_PREFIX = "gs://"
|
|
_S3_PREFIX = "s3://"
|
|
_HDFS_PREFIX = "hdfs://"
|
|
_WEBHDFS_PREFIX = "webhdfs://"
|
|
_AZURE_BLOB_RE = [
|
|
"https://(.+?).blob.core.windows.net/(.+)",
|
|
"https://(.+?).z[0-9]{1,2}.blob.storage.azure.net/(.+)",
|
|
]
|
|
_AZURE_FILE_RE = [
|
|
"https://(.+?).file.core.windows.net/(.+)",
|
|
"https://(.+?).z[0-9]{1,2}.file.storage.azure.net/(.+)",
|
|
]
|
|
_LOCAL_PREFIX = "file://"
|
|
_URI_RE = "https?://(.+)/(.+)"
|
|
_HTTP_PREFIX = "http(s)://"
|
|
_HEADERS_SUFFIX = "-headers"
|
|
_PVC_PREFIX = "/mnt/pvc"
|
|
_HF_PREFIX = "hf://"
|
|
_GIT_RE = r"https://.+\.git"
|
|
|
|
_HDFS_SECRET_DIRECTORY = "/var/secrets/kserve-hdfscreds"
|
|
_HDFS_FILE_SECRETS = ["KERBEROS_KEYTAB", "TLS_CERT", "TLS_KEY", "TLS_CA"]
|
|
|
|
# S3 parallel download configuration
|
|
_S3_MAX_FILE_CONCURRENCY = int(os.getenv("S3_MAX_FILE_CONCURRENCY", "4"))
|
|
# Global variable for S3 resource in worker processes
|
|
_worker_s3_resource = None
|
|
# Azure async download configuration
|
|
_AZURE_MAX_FILE_CONCURRENCY = int(os.getenv("AZURE_MAX_FILE_CONCURRENCY", "4"))
|
|
_AZURE_MAX_CHUNK_CONCURRENCY = int(os.getenv("AZURE_MAX_CHUNK_CONCURRENCY", "4"))
|
|
|
|
|
|
def _should_download(
|
|
relative_path: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> bool:
|
|
"""Determine whether a file should be downloaded based on allow/ignore patterns.
|
|
|
|
Uses fnmatch semantics consistent with huggingface_hub.snapshot_download.
|
|
Matching is performed against the relative path of the file.
|
|
|
|
Args:
|
|
relative_path: The relative path of the file (e.g., "subdir/model.safetensors").
|
|
allow_patterns: If set, only files matching at least one pattern are included.
|
|
ignore_patterns: If set, files matching any pattern are excluded.
|
|
|
|
Returns:
|
|
True if the file should be downloaded, False otherwise.
|
|
"""
|
|
if allow_patterns is not None and len(allow_patterns) > 0:
|
|
if not any(fnmatch.fnmatch(relative_path, p) for p in allow_patterns):
|
|
return False
|
|
|
|
if ignore_patterns is not None and len(ignore_patterns) > 0:
|
|
if any(fnmatch.fnmatch(relative_path, p) for p in ignore_patterns):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def _parse_patterns_from_env(env_var_name: str) -> Optional[List[str]]:
|
|
"""Parse allow/ignore patterns from an environment variable.
|
|
|
|
Supports JSON array format: '["*.safetensors", "*.json"]'
|
|
Falls back to comma-separated: '*.safetensors,*.json'
|
|
|
|
Returns None if the env var is not set or empty.
|
|
"""
|
|
value = os.environ.get(env_var_name, "").strip()
|
|
if not value:
|
|
return None
|
|
|
|
try:
|
|
parsed = json.loads(value)
|
|
if isinstance(parsed, list):
|
|
return [str(p).strip() for p in parsed if str(p).strip()]
|
|
return [str(parsed).strip()] if str(parsed).strip() else None
|
|
except (json.JSONDecodeError, TypeError):
|
|
patterns = [p.strip() for p in value.split(",") if p.strip()]
|
|
return patterns if patterns else None
|
|
|
|
|
|
class Storage(object):
|
|
@staticmethod
|
|
def download_files(
|
|
source_uris: list[str],
|
|
out_dirs: list[str],
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> list[str]:
|
|
for d in out_dirs:
|
|
if d:
|
|
os.makedirs(d, exist_ok=True)
|
|
download_fn = partial(
|
|
Storage.download,
|
|
allow_patterns=allow_patterns,
|
|
ignore_patterns=ignore_patterns,
|
|
)
|
|
# Sequential: parallel snapshot_download of the same Hugging Face repo into
|
|
# different local_dir paths races on the process-wide HF cache (locks, temp files).
|
|
return [
|
|
download_fn(uri, out)
|
|
for uri, out in zip(source_uris, out_dirs, strict=True)
|
|
]
|
|
|
|
@staticmethod
|
|
def download(
|
|
uri: str,
|
|
out_dir: Optional[str] = None,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
start = time.monotonic()
|
|
Storage._update_with_storage_spec()
|
|
|
|
if allow_patterns is None:
|
|
allow_patterns = _parse_patterns_from_env("STORAGE_ALLOW_PATTERNS")
|
|
if ignore_patterns is None:
|
|
ignore_patterns = _parse_patterns_from_env("STORAGE_IGNORE_PATTERNS")
|
|
|
|
logger.info("Copying contents of %s to local", uri)
|
|
|
|
if allow_patterns:
|
|
logger.info("Allow patterns: %s", allow_patterns)
|
|
if ignore_patterns:
|
|
logger.info("Ignore patterns: %s", ignore_patterns)
|
|
|
|
if uri.startswith(_PVC_PREFIX) and not os.path.exists(uri):
|
|
raise Exception(f"Cannot locate source uri {uri} for PVC")
|
|
|
|
is_local = uri.startswith(_LOCAL_PREFIX) or os.path.exists(uri)
|
|
if is_local:
|
|
if out_dir is None:
|
|
# noop if out_dir is not set and the path is local
|
|
model_dir = Storage._download_local(uri)
|
|
else:
|
|
if not os.path.exists(out_dir):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
model_dir = Storage._download_local(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
else:
|
|
if out_dir is None:
|
|
out_dir = tempfile.mkdtemp()
|
|
elif not os.path.exists(out_dir):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
if uri.startswith(MODEL_MOUNT_DIRS):
|
|
# Don't need to download models if this InferenceService is running in the multi-model
|
|
# serving mode. The model agent will download models.
|
|
model_dir = out_dir
|
|
elif uri.startswith(_GCS_PREFIX):
|
|
model_dir = Storage._download_gcs(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif uri.startswith(_S3_PREFIX):
|
|
model_dir = Storage._download_s3(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif uri.startswith(_HDFS_PREFIX) or uri.startswith(_WEBHDFS_PREFIX):
|
|
model_dir = Storage._download_hdfs(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif any(re.search(pattern, uri) for pattern in _AZURE_BLOB_RE):
|
|
model_dir = Storage._download_azure_blob(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif any(re.search(pattern, uri) for pattern in _AZURE_FILE_RE):
|
|
model_dir = Storage._download_azure_file_share(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif uri.startswith(_HF_PREFIX):
|
|
model_dir = Storage._download_hf(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
elif re.search(_GIT_RE, uri):
|
|
model_dir = Storage._download_git_repo(uri, out_dir)
|
|
# "catch-all" pattern, should always be last
|
|
elif re.search(_URI_RE, uri):
|
|
model_dir = Storage._download_from_uri(uri, out_dir)
|
|
else:
|
|
raise Exception(
|
|
"Cannot recognize storage type for "
|
|
+ uri
|
|
+ "\n'%s', '%s', '%s', '%s' and '%s' are the current available storage type."
|
|
% (_GCS_PREFIX, _S3_PREFIX, _LOCAL_PREFIX, _HTTP_PREFIX, _HF_PREFIX)
|
|
)
|
|
|
|
logger.info("Successfully copied %s to %s", uri, out_dir)
|
|
logger.info(f"Model downloaded in {time.monotonic() - start} seconds.")
|
|
return model_dir
|
|
|
|
@staticmethod
|
|
def _update_with_storage_spec():
|
|
storage_secret_json = json.loads(os.environ.get("STORAGE_CONFIG", "{}"))
|
|
storage_secret_override_params = json.loads(
|
|
os.environ.get("STORAGE_OVERRIDE_CONFIG", "{}")
|
|
)
|
|
if storage_secret_override_params:
|
|
for key, value in storage_secret_override_params.items():
|
|
storage_secret_json[key] = value
|
|
|
|
if storage_secret_json.get("type", "") == "s3":
|
|
for env_var, key in (
|
|
("AWS_ENDPOINT_URL", "endpoint_url"),
|
|
("AWS_ACCESS_KEY_ID", "access_key_id"),
|
|
("AWS_SECRET_ACCESS_KEY", "secret_access_key"),
|
|
("AWS_DEFAULT_REGION", "region"),
|
|
("AWS_CA_BUNDLE", "ca_bundle"),
|
|
("S3_VERIFY_SSL", "verify_ssl"),
|
|
("awsAnonymousCredential", "anonymous"),
|
|
):
|
|
if key in storage_secret_json:
|
|
os.environ[env_var] = storage_secret_json.get(key)
|
|
|
|
if (
|
|
storage_secret_json.get("type", "") == "hdfs"
|
|
or storage_secret_json.get("type", "") == "webhdfs"
|
|
):
|
|
temp_dir = tempfile.mkdtemp()
|
|
os.environ["HDFS_SECRET_DIR"] = temp_dir
|
|
for key, value in storage_secret_json.items():
|
|
mode = "w"
|
|
|
|
# If the secret is supposed to be a file, then it was base64 encoded in the json
|
|
if key in _HDFS_FILE_SECRETS:
|
|
value = base64.b64decode(value)
|
|
mode = "wb"
|
|
|
|
with open(f"{temp_dir}/{key}", mode) as f:
|
|
f.write(value)
|
|
f.flush()
|
|
|
|
@staticmethod
|
|
def get_S3_config():
|
|
from botocore import UNSIGNED
|
|
from botocore.client import Config
|
|
|
|
# Configurable timeouts and retries
|
|
connect_timeout = int(os.getenv("S3_CONNECT_TIMEOUT", "15"))
|
|
read_timeout = int(os.getenv("S3_READ_TIMEOUT", "30"))
|
|
max_attempts = int(os.getenv("S3_MAX_ATTEMPTS", "3"))
|
|
|
|
# default s3 config with timeouts
|
|
c = Config(
|
|
connect_timeout=connect_timeout,
|
|
read_timeout=read_timeout,
|
|
retries={"max_attempts": max_attempts, "mode": "standard"},
|
|
)
|
|
|
|
# anon environment variable defined in s3_secret.go
|
|
anon = "true" == os.getenv("awsAnonymousCredential", "false").lower()
|
|
# S3UseVirtualBucket environment variable defined in s3_secret.go
|
|
# use virtual hosted-style URLs if enabled
|
|
virtual = "true" == os.getenv("S3_USER_VIRTUAL_BUCKET", "false").lower()
|
|
# S3UseAccelerate environment variable defined in s3_secret.go
|
|
# use transfer acceleration if enabled
|
|
accelerate = "true" == os.getenv("S3_USE_ACCELERATE", "false").lower()
|
|
|
|
if anon:
|
|
c = c.merge(Config(signature_version=UNSIGNED))
|
|
if virtual:
|
|
c = c.merge(Config(s3={"addressing_style": "virtual"}))
|
|
if accelerate:
|
|
c = c.merge(Config(s3={"use_accelerate_endpoint": accelerate}))
|
|
|
|
# NOTE: If endpoint_url provided is legacy ("https://s3.amazonaws.com") and region is not global (us-east-1), set to virtual addressing style
|
|
# So that request would not return PermanentRedirect due to region not in the endpoint url
|
|
# AWS SDK retries under the hood to set the correct region when the valid virtual addressing style endpoint url is provided
|
|
endpoint_url = os.getenv("AWS_ENDPOINT_URL")
|
|
region = os.getenv("AWS_DEFAULT_REGION")
|
|
if endpoint_url == "https://s3.amazonaws.com" and region not in (
|
|
None,
|
|
"us-east-1",
|
|
):
|
|
c = c.merge(Config(s3={"addressing_style": "virtual"}))
|
|
|
|
return c
|
|
|
|
@staticmethod
|
|
def _get_s3_client_kwargs():
|
|
"""
|
|
Get the standardized boto3 client kwargs for S3 operations.
|
|
|
|
Returns:
|
|
dict: kwargs for creating boto3 S3 resource/client
|
|
"""
|
|
kwargs = {"config": Storage.get_S3_config()}
|
|
endpoint_url = os.getenv("AWS_ENDPOINT_URL")
|
|
if endpoint_url:
|
|
kwargs.update({"endpoint_url": endpoint_url})
|
|
verify_ssl = os.getenv("S3_VERIFY_SSL")
|
|
if verify_ssl:
|
|
verify_ssl = verify_ssl.lower() not in ["0", "false"]
|
|
kwargs.update({"verify": verify_ssl})
|
|
else:
|
|
verify_ssl = True
|
|
|
|
# If verify_ssl is true, then check there is custom ca bundle cert
|
|
# The CA bundle can be any local file in the container under the path
|
|
# set in the AWS_CA_BUNDLE environment variable.
|
|
# It can also be coming from a ConfigMap, in which case the filename
|
|
# is cabundle.crt.
|
|
if verify_ssl:
|
|
global_ca_bundle_configmap = os.getenv("CA_BUNDLE_CONFIGMAP_NAME")
|
|
isvc_aws_ca_bundle_path = os.getenv("AWS_CA_BUNDLE")
|
|
ca_bundle_set = False
|
|
if isvc_aws_ca_bundle_path and isvc_aws_ca_bundle_path != "":
|
|
ca_bundle_set = True
|
|
ca_bundle_full_path = isvc_aws_ca_bundle_path
|
|
elif global_ca_bundle_configmap:
|
|
ca_bundle_set = True
|
|
global_ca_bundle_volume_mount_path = os.getenv(
|
|
"CA_BUNDLE_VOLUME_MOUNT_POINT"
|
|
)
|
|
ca_bundle_full_path = os.path.join(
|
|
global_ca_bundle_volume_mount_path, "cabundle.crt"
|
|
)
|
|
if ca_bundle_set:
|
|
if os.path.exists(ca_bundle_full_path):
|
|
logger.info("ca bundle file(%s) exists." % (ca_bundle_full_path))
|
|
kwargs.update({"verify": ca_bundle_full_path})
|
|
else:
|
|
raise RuntimeError(
|
|
"Failed to find ca bundle file(%s)." % ca_bundle_full_path
|
|
)
|
|
return kwargs
|
|
|
|
@staticmethod
|
|
def _init_s3_worker():
|
|
"""
|
|
Initialize S3 resources for worker processes.
|
|
This runs once per worker process to avoid repeated initialization overhead.
|
|
"""
|
|
global _worker_s3_resource
|
|
try:
|
|
import boto3
|
|
|
|
kwargs = Storage._get_s3_client_kwargs()
|
|
session = boto3.Session()
|
|
_worker_s3_resource = session.resource("s3", **kwargs)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize S3 worker: {e}")
|
|
_worker_s3_resource = None
|
|
|
|
@staticmethod
|
|
def _download_s3_object(args: Tuple) -> Tuple[bool, str, str]:
|
|
"""
|
|
Worker function to download a single S3 object in a separate process.
|
|
Uses pre-initialized S3 resource from _init_s3_worker().
|
|
|
|
Args:
|
|
args: Tuple containing (bucket_name:str, obj_key:str, target_path:str)
|
|
|
|
Returns:
|
|
Tuple of (success: bool, obj_key: str, error_message: str)
|
|
"""
|
|
try:
|
|
if _worker_s3_resource is None:
|
|
return False, args[1], "S3 resource not initialized in worker process"
|
|
|
|
bucket_name, obj_key, target_path = args
|
|
bucket = _worker_s3_resource.Bucket(bucket_name)
|
|
|
|
# Download the file
|
|
bucket.download_file(obj_key, target_path)
|
|
|
|
return True, obj_key, ""
|
|
|
|
except Exception as e:
|
|
return False, obj_key, str(e)
|
|
|
|
@staticmethod
|
|
def _download_s3(
|
|
uri,
|
|
temp_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
import boto3
|
|
from botocore.exceptions import ClientError, NoCredentialsError
|
|
|
|
# Get S3 configuration using the shared helper method
|
|
kwargs = Storage._get_s3_client_kwargs()
|
|
parsed = urlparse(uri, scheme="s3")
|
|
bucket_name = parsed.netloc
|
|
bucket_path = parsed.path.lstrip("/")
|
|
|
|
try:
|
|
s3 = boto3.resource("s3", **kwargs)
|
|
except ValueError as e:
|
|
endpoint_url = os.getenv("AWS_ENDPOINT_URL", "")
|
|
raise_storage_error("S3", uri, e, endpoint_url)
|
|
|
|
# Collect all objects to download
|
|
download_tasks = []
|
|
exact_obj_found = False
|
|
|
|
try:
|
|
s3.meta.client.head_bucket(Bucket=bucket_name)
|
|
bucket = s3.Bucket(bucket_name)
|
|
|
|
for obj in bucket.objects.filter(Prefix=bucket_path):
|
|
if obj.key.endswith("/") or obj.size == 0:
|
|
logger.debug("Skipping: %s", obj.key)
|
|
continue
|
|
|
|
logger.info("Found S3 object: %s (%d bytes)", obj.key, obj.size)
|
|
|
|
if bucket_path == obj.key:
|
|
target_key = obj.key.rsplit("/", 1)[-1]
|
|
exact_obj_found = True
|
|
else:
|
|
target_key = obj.key.removeprefix(bucket_path).lstrip("/")
|
|
|
|
# Apply file filtering (skip for exact object match)
|
|
if not exact_obj_found and not _should_download(
|
|
target_key, allow_patterns, ignore_patterns
|
|
):
|
|
logger.info("Skipping %s due to file pattern filter", obj.key)
|
|
continue
|
|
|
|
target_path = f"{temp_dir}/{target_key}"
|
|
|
|
# Create target directory if it doesn't exist
|
|
if not os.path.exists(dir_path := os.path.dirname(target_path)):
|
|
os.makedirs(dir_path, exist_ok=True)
|
|
|
|
download_tasks.append((bucket_name, obj.key, target_path))
|
|
|
|
# If the exact object is found, then it is sufficient to download that and break the loop
|
|
if exact_obj_found:
|
|
break
|
|
except (ClientError, NoCredentialsError) as e:
|
|
raise_storage_error("S3", uri, e, bucket_name)
|
|
|
|
if len(download_tasks) == 0:
|
|
raise RuntimeError(
|
|
"Failed to fetch model. No model found in %s." % bucket_path
|
|
)
|
|
|
|
num_processes = min(_S3_MAX_FILE_CONCURRENCY, len(download_tasks))
|
|
|
|
with multiprocessing.Pool(
|
|
processes=num_processes, initializer=Storage._init_s3_worker
|
|
) as pool:
|
|
results = list(
|
|
pool.imap_unordered(
|
|
Storage._download_s3_object, download_tasks, chunksize=1
|
|
)
|
|
)
|
|
|
|
# Process results and handle errors
|
|
successful_downloads = []
|
|
failed_downloads = []
|
|
|
|
for success, obj_key, error_msg in results:
|
|
if success:
|
|
successful_downloads.append(obj_key)
|
|
logger.info("Downloaded object %s" % obj_key)
|
|
else:
|
|
failed_downloads.append((obj_key, error_msg))
|
|
logger.error("Failed to download object %s: %s" % (obj_key, error_msg))
|
|
|
|
if len(failed_downloads) > 0:
|
|
error_details = "; ".join(
|
|
[f"{obj}: {err}" for obj, err in failed_downloads]
|
|
)
|
|
raise RuntimeError(
|
|
f"Failed to download {len(failed_downloads)} files: {error_details}"
|
|
)
|
|
|
|
# Unpack compressed file, supports .tgz, tar.gz and zip file formats.
|
|
if len(successful_downloads) == 1:
|
|
target_path = download_tasks[0][
|
|
2
|
|
] # target_path is the 3rd element in task_args
|
|
mimetype, _ = mimetypes.guess_type(target_path)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
temp_dir = Storage._unpack_archive_file(target_path, mimetype, temp_dir)
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def _download_hf(
|
|
uri,
|
|
temp_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
from huggingface_hub import snapshot_download
|
|
|
|
from huggingface_hub.utils import (
|
|
RepositoryNotFoundError,
|
|
RevisionNotFoundError,
|
|
GatedRepoError,
|
|
HfHubHTTPError,
|
|
)
|
|
|
|
components = uri[len(_HF_PREFIX) :].split("/")
|
|
|
|
# Validate that the URI has two parts: repo and model (optional hash)
|
|
if len(components) != 2:
|
|
raise RuntimeError(
|
|
"Invalid Hugging Face URI format. Expected 'hf://owner/model[:revision]', got '%s'"
|
|
% uri
|
|
)
|
|
|
|
repo = components[0]
|
|
model_part = components[1]
|
|
|
|
if not repo:
|
|
raise RuntimeError(
|
|
"Hugging Face repository owner cannot be empty in URI: %s" % uri
|
|
)
|
|
if not model_part:
|
|
raise RuntimeError(
|
|
"Hugging Face model name cannot be empty in URI: %s" % uri
|
|
)
|
|
|
|
model, _, hash_value = model_part.partition(":")
|
|
# Ensure model is non-empty
|
|
if not model:
|
|
raise RuntimeError(
|
|
"Hugging Face model name cannot be empty in URI: %s" % uri
|
|
)
|
|
|
|
revision = hash_value if hash_value else None
|
|
repo_id = f"{repo}/{model}"
|
|
|
|
try:
|
|
kwargs = dict(repo_id=repo_id, revision=revision, local_dir=temp_dir)
|
|
if allow_patterns:
|
|
kwargs["allow_patterns"] = allow_patterns
|
|
if ignore_patterns:
|
|
kwargs["ignore_patterns"] = ignore_patterns
|
|
snapshot_download(**kwargs)
|
|
except (
|
|
RepositoryNotFoundError,
|
|
RevisionNotFoundError,
|
|
GatedRepoError,
|
|
HfHubHTTPError,
|
|
) as e:
|
|
raise_storage_error("HuggingFace", uri, e, repo_id)
|
|
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def _download_gcs(
|
|
uri,
|
|
temp_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
from google.auth import exceptions as auth_exceptions
|
|
from google.cloud import storage
|
|
from google.api_core import exceptions as api_exceptions
|
|
import copy
|
|
|
|
try:
|
|
storage_client = storage.Client()
|
|
except auth_exceptions.DefaultCredentialsError:
|
|
storage_client = storage.Client.create_anonymous_client()
|
|
|
|
bucket_args = uri.replace(_GCS_PREFIX, "", 1).split("/", 1)
|
|
bucket_name = bucket_args[0]
|
|
bucket_path = bucket_args[1] if len(bucket_args) > 1 else ""
|
|
|
|
try:
|
|
bucket = storage_client.bucket(bucket_name)
|
|
prefix = bucket_path
|
|
if not prefix.endswith("/"):
|
|
prefix = prefix + "/"
|
|
blobs = bucket.list_blobs(prefix=prefix)
|
|
file_count = 0
|
|
|
|
# Shallow copy, otherwise Iterator has already started
|
|
shallow_blobs = copy.copy(blobs)
|
|
blob = bucket.blob(bucket_path)
|
|
# checks if the blob is a file or a directory
|
|
if blob.name == bucket_path and len(list(shallow_blobs)) == 0:
|
|
dest_path = os.path.join(temp_dir, os.path.basename(bucket_path))
|
|
logger.info("Downloading single file to: %s", dest_path)
|
|
blob.download_to_filename(dest_path)
|
|
file_count = 1
|
|
|
|
else:
|
|
for blob in blobs:
|
|
# Replace any prefix from the object key with temp_dir
|
|
subdir_object_key = blob.name.replace(bucket_path, "", 1).lstrip(
|
|
"/"
|
|
)
|
|
# Create necessary subdirectory to store the object locally
|
|
if "/" in subdir_object_key:
|
|
local_object_dir = os.path.join(
|
|
temp_dir, subdir_object_key.rsplit("/", 1)[0]
|
|
)
|
|
if not os.path.isdir(local_object_dir):
|
|
os.makedirs(local_object_dir, exist_ok=True)
|
|
if (
|
|
subdir_object_key.strip() != ""
|
|
and not subdir_object_key.endswith("/")
|
|
):
|
|
if not _should_download(
|
|
subdir_object_key, allow_patterns, ignore_patterns
|
|
):
|
|
logger.info(
|
|
"Skipping %s due to file pattern filter", blob.name
|
|
)
|
|
continue
|
|
dest_path = os.path.join(temp_dir, subdir_object_key)
|
|
logger.info("Downloading: %s", dest_path)
|
|
blob.download_to_filename(dest_path)
|
|
file_count += 1
|
|
except (
|
|
api_exceptions.Forbidden,
|
|
api_exceptions.NotFound,
|
|
auth_exceptions.GoogleAuthError,
|
|
) as e:
|
|
raise_storage_error("GCS", uri, e, bucket_name)
|
|
|
|
if file_count == 0:
|
|
raise RuntimeError("Failed to fetch model. No model found in %s." % uri)
|
|
|
|
# Unpack compressed file, supports .tgz, tar.gz and zip file formats.
|
|
if file_count == 1:
|
|
mimetype, _ = mimetypes.guess_type(blob.name)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
temp_dir = Storage._unpack_archive_file(dest_path, mimetype, temp_dir)
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def _load_hdfs_configuration() -> dict:
|
|
config = {
|
|
"HDFS_NAMENODE": None,
|
|
"USER_PROXY": None,
|
|
"HDFS_ROOTPATH": None,
|
|
"TLS_CERT": None,
|
|
"TLS_KEY": None,
|
|
"TLS_CA": None,
|
|
"TLS_SKIP_VERIFY": "false",
|
|
"HEADERS": None,
|
|
"N_THREADS": "2",
|
|
"KERBEROS_KEYTAB": None,
|
|
"KERBEROS_PRINCIPAL": None,
|
|
}
|
|
|
|
secret_dir = _HDFS_SECRET_DIRECTORY
|
|
if os.environ.get("HDFS_SECRET_DIR"):
|
|
secret_dir = os.environ["HDFS_SECRET_DIR"]
|
|
|
|
for filename in os.listdir(secret_dir):
|
|
if filename not in config:
|
|
continue
|
|
|
|
# We don't read files which are supposed to be files, just save their path
|
|
if filename in _HDFS_FILE_SECRETS:
|
|
config[filename] = f"{secret_dir}/{filename}"
|
|
continue
|
|
|
|
# Read file and save value in config dict
|
|
with open(f"{secret_dir}/{filename}") as f:
|
|
config[filename] = f.read()
|
|
|
|
return config
|
|
|
|
@staticmethod
|
|
def _download_hdfs(
|
|
uri,
|
|
out_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
from krbcontext.context import krbContext
|
|
from hdfs.ext.kerberos import Client, KerberosClient
|
|
from hdfs.util import HdfsError
|
|
|
|
config = Storage._load_hdfs_configuration()
|
|
|
|
logger.info("Using the following hdfs config: %s", config)
|
|
|
|
# Remove hdfs:// or webhdfs:// from the uri to get just the path
|
|
# e.g. hdfs://user/me/model -> user/me/model
|
|
if uri.startswith(_HDFS_PREFIX):
|
|
path = uri[len(_HDFS_PREFIX) :]
|
|
else:
|
|
path = uri[len(_WEBHDFS_PREFIX) :]
|
|
|
|
if not config["HDFS_ROOTPATH"]:
|
|
path = "/" + path
|
|
|
|
s = requests.Session()
|
|
|
|
if config["TLS_CERT"]:
|
|
s.cert = (config["TLS_CERT"], config["TLS_KEY"])
|
|
# s.verify = , True, False, or CA PATH
|
|
if config["TLS_CA"]:
|
|
s.verify = config["TLS_CA"]
|
|
if config["TLS_SKIP_VERIFY"].lower() == "true":
|
|
s.verify = False
|
|
|
|
if config["HEADERS"]:
|
|
headers = json.loads(config["HEADERS"])
|
|
s.headers.update(headers)
|
|
|
|
try:
|
|
if config["KERBEROS_PRINCIPAL"]:
|
|
context = krbContext(
|
|
using_keytab=True,
|
|
principal=config["KERBEROS_PRINCIPAL"],
|
|
keytab_file=config["KERBEROS_KEYTAB"],
|
|
)
|
|
context.init_with_keytab()
|
|
client = KerberosClient(
|
|
config["HDFS_NAMENODE"],
|
|
proxy=config["USER_PROXY"],
|
|
root=config["HDFS_ROOTPATH"],
|
|
session=s,
|
|
)
|
|
else:
|
|
client = Client(
|
|
config["HDFS_NAMENODE"],
|
|
proxy=config["USER_PROXY"],
|
|
root=config["HDFS_ROOTPATH"],
|
|
session=s,
|
|
)
|
|
file_count = 0
|
|
dest_file_path = ""
|
|
|
|
# Check path exists and get path status
|
|
# Raises HdfsError when path does not exist
|
|
status = client.status(path)
|
|
|
|
if status["type"] == "FILE":
|
|
client.download(path, out_dir, n_threads=1)
|
|
file_count += 1
|
|
file_name = path.rsplit("/", 1)[-1]
|
|
dest_file_path = f"{out_dir}/{file_name}"
|
|
else:
|
|
files = client.list(path)
|
|
for f in files:
|
|
if not _should_download(f, allow_patterns, ignore_patterns):
|
|
logger.info("Skipping %s due to file pattern filter", f)
|
|
continue
|
|
client.download(
|
|
f"{path}/{f}", out_dir, n_threads=int(config["N_THREADS"])
|
|
)
|
|
file_count += 1
|
|
dest_file_path = f"{out_dir}/{f}"
|
|
except (HdfsError, requests.exceptions.ConnectionError) as e:
|
|
raise_storage_error("HDFS", uri, e, config["HDFS_NAMENODE"])
|
|
|
|
if file_count == 1:
|
|
mimetype, _ = mimetypes.guess_type(dest_file_path)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
out_dir = Storage._unpack_archive_file(
|
|
dest_file_path, mimetype, out_dir
|
|
)
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
async def _download_azure_blob_async(
|
|
uri,
|
|
out_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
"""Async Azure blob download with chunked streaming and multi-level semaphores"""
|
|
from azure.storage.blob.aio import BlobServiceClient
|
|
from azure.core.exceptions import (
|
|
ClientAuthenticationError,
|
|
ResourceNotFoundError,
|
|
HttpResponseError,
|
|
)
|
|
|
|
account_name, account_url, container_name, prefix = Storage._parse_azure_uri(
|
|
uri
|
|
)
|
|
logger.info(
|
|
"Connecting to BLOB account: [%s], container: [%s], prefix: [%s]",
|
|
account_name,
|
|
container_name,
|
|
prefix,
|
|
)
|
|
|
|
token = (
|
|
Storage._get_azure_storage_token()
|
|
or Storage._get_azure_storage_access_key()
|
|
)
|
|
if token is None:
|
|
logger.warning(
|
|
"Azure credentials or shared access signature token not found, retrying anonymous access"
|
|
)
|
|
|
|
# File-level semaphore to control concurrent file downloads
|
|
file_semaphore = asyncio.Semaphore(_AZURE_MAX_FILE_CONCURRENCY)
|
|
|
|
try:
|
|
async with BlobServiceClient(
|
|
account_url, credential=token
|
|
) as blob_service_client:
|
|
container_client = blob_service_client.get_container_client(
|
|
container_name
|
|
)
|
|
|
|
# Get all blobs using flat listing (no delimiter) to get all files regardless of hierarchy
|
|
blobs = []
|
|
logger.info("Listing blobs with prefix: %s", prefix)
|
|
async for blob in container_client.list_blobs(name_starts_with=prefix):
|
|
logger.info("Found blob: %s (%d bytes)", blob.name, blob.size)
|
|
if blob.size > 0:
|
|
relative_path = blob.name.replace(prefix, "", 1).lstrip("/")
|
|
if not relative_path:
|
|
relative_path = os.path.basename(prefix)
|
|
if not _should_download(
|
|
relative_path, allow_patterns, ignore_patterns
|
|
):
|
|
logger.info(
|
|
"Skipping %s due to file pattern filter", blob.name
|
|
)
|
|
continue
|
|
blobs.append(blob)
|
|
|
|
if not blobs:
|
|
raise RuntimeError(
|
|
"Failed to fetch model. No model found in %s." % uri
|
|
)
|
|
|
|
# Create download tasks with semaphore control
|
|
download_tasks = [
|
|
Storage._download_single_blob_async(
|
|
container_client, blob, out_dir, prefix, file_semaphore
|
|
)
|
|
for blob in blobs
|
|
]
|
|
|
|
# Execute all downloads concurrently
|
|
results = await asyncio.gather(*download_tasks, return_exceptions=True)
|
|
|
|
# Check for exceptions
|
|
file_count = 0
|
|
dest_path = None
|
|
for i, result in enumerate(results):
|
|
if isinstance(result, Exception):
|
|
logger.error("Blob %s failed: %s", blobs[i].name, result)
|
|
raise result
|
|
else:
|
|
dest_path = result
|
|
file_count += 1
|
|
except (
|
|
ClientAuthenticationError,
|
|
ResourceNotFoundError,
|
|
HttpResponseError,
|
|
) as e:
|
|
raise_storage_error("Azure", uri, e, container_name)
|
|
|
|
# Handle single file unpacking
|
|
if file_count == 1:
|
|
mimetype, _ = mimetypes.guess_type(dest_path)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)
|
|
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
async def _download_single_blob_async(
|
|
container_client,
|
|
blob,
|
|
out_dir: str,
|
|
prefix: str,
|
|
file_semaphore: asyncio.Semaphore,
|
|
) -> str:
|
|
"""Download a single blob with chunked streaming"""
|
|
async with file_semaphore: # Limit concurrent file downloads
|
|
file_name = blob.name.replace(prefix, "", 1).lstrip("/")
|
|
if not file_name:
|
|
file_name = os.path.basename(prefix)
|
|
dest_path = os.path.join(out_dir, file_name)
|
|
Path(os.path.dirname(dest_path)).mkdir(parents=True, exist_ok=True)
|
|
|
|
logger.info("Downloading: %s to %s", blob.name, dest_path)
|
|
|
|
# Get blob client and download with chunked streaming
|
|
blob_client = container_client.get_blob_client(blob.name)
|
|
|
|
# Download with streaming chunks to avoid memory overload
|
|
stream = await blob_client.download_blob(
|
|
max_concurrency=_AZURE_MAX_CHUNK_CONCURRENCY
|
|
)
|
|
|
|
with open(dest_path, "wb") as f:
|
|
# Stream chunks instead of loading entire file into memory
|
|
async for chunk in stream.chunks():
|
|
f.write(chunk)
|
|
|
|
return dest_path
|
|
|
|
@staticmethod
|
|
def _download_azure_file_share(
|
|
uri,
|
|
out_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str: # pylint: disable=too-many-locals
|
|
from azure.storage.fileshare import ShareServiceClient
|
|
from azure.core.exceptions import (
|
|
ClientAuthenticationError,
|
|
ResourceNotFoundError,
|
|
HttpResponseError,
|
|
)
|
|
|
|
account_name, account_url, share_name, prefix = Storage._parse_azure_uri(uri)
|
|
logger.info(
|
|
"Connecting to file share account: [%s], container: [%s], prefix: [%s]",
|
|
account_name,
|
|
share_name,
|
|
prefix,
|
|
)
|
|
access_key = Storage._get_azure_storage_access_key()
|
|
if access_key is None:
|
|
logger.warning(
|
|
"Azure storage access key not found, retrying anonymous access"
|
|
)
|
|
|
|
try:
|
|
share_service_client = ShareServiceClient(
|
|
account_url, credential=access_key
|
|
)
|
|
share_client = share_service_client.get_share_client(share_name)
|
|
file_count = 0
|
|
share_files = []
|
|
max_depth = 5
|
|
stack = [(prefix, max_depth)]
|
|
while stack:
|
|
curr_prefix, depth = stack.pop()
|
|
if depth < 0:
|
|
continue
|
|
for item in share_client.list_directories_and_files(
|
|
directory_name=curr_prefix
|
|
):
|
|
if item.is_directory:
|
|
stack.append(
|
|
("/".join([curr_prefix, item.name]).strip("/"), depth - 1)
|
|
)
|
|
else:
|
|
share_files.append((curr_prefix, item))
|
|
for prefix, file_item in share_files:
|
|
parts = [prefix] if prefix else []
|
|
parts.append(file_item.name)
|
|
file_path = "/".join(parts).lstrip("/")
|
|
if not _should_download(file_path, allow_patterns, ignore_patterns):
|
|
logger.info("Skipping %s due to file pattern filter", file_path)
|
|
continue
|
|
dest_path = os.path.join(out_dir, file_path)
|
|
Path(os.path.dirname(dest_path)).mkdir(parents=True, exist_ok=True)
|
|
logger.info("Downloading: %s to %s", file_item.name, dest_path)
|
|
file_client = share_client.get_file_client(file_path)
|
|
with open(dest_path, "wb+") as f:
|
|
data = file_client.download_file()
|
|
data.readinto(f)
|
|
file_count += 1
|
|
except (
|
|
ClientAuthenticationError,
|
|
ResourceNotFoundError,
|
|
HttpResponseError,
|
|
) as e:
|
|
raise_storage_error("Azure", uri, e, share_name)
|
|
|
|
if file_count == 0:
|
|
raise RuntimeError("Failed to fetch model. No model found in %s." % (uri))
|
|
|
|
# Unpack compressed file, supports .tgz, tar.gz and zip file formats.
|
|
if file_count == 1:
|
|
mimetype, _ = mimetypes.guess_type(dest_path)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
def _download_azure_blob(
|
|
uri,
|
|
out_dir: str,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
"""Wrapper to run async blob download"""
|
|
return asyncio.run(
|
|
Storage._download_azure_blob_async(
|
|
uri, out_dir, allow_patterns, ignore_patterns
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _parse_azure_uri(uri): # pylint: disable=too-many-locals
|
|
parsed = urlparse(uri)
|
|
account_name = parsed.netloc.split(".")[0]
|
|
account_url = "https://{}{}".format(
|
|
parsed.netloc, "?" + parsed.query if parsed.query else ""
|
|
)
|
|
object_name, prefix = parsed.path.lstrip("/").split("/", 1)
|
|
prefix = prefix.strip("/")
|
|
return account_name, account_url, object_name, prefix
|
|
|
|
@staticmethod
|
|
def _get_azure_storage_token():
|
|
tenant_id = os.getenv("AZ_TENANT_ID", "")
|
|
client_id = os.getenv("AZ_CLIENT_ID", "")
|
|
client_secret = os.getenv("AZ_CLIENT_SECRET", "")
|
|
|
|
# convert old environment variable to conform Azure defaults
|
|
# see azure/identity/_constants.py
|
|
if tenant_id:
|
|
os.environ["AZURE_TENANT_ID"] = tenant_id
|
|
if client_id:
|
|
os.environ["AZURE_CLIENT_ID"] = client_id
|
|
if client_secret:
|
|
os.environ["AZURE_CLIENT_SECRET"] = client_secret
|
|
|
|
client_id = os.getenv("AZURE_CLIENT_ID", "")
|
|
if not client_id:
|
|
return None
|
|
|
|
# note the SP must have "Storage Blob Data Owner" perms for this to work
|
|
from azure.identity import DefaultAzureCredential
|
|
|
|
token_credential = DefaultAzureCredential()
|
|
|
|
logger.info("Retrieved SP token credential for client_id: %s", client_id)
|
|
return token_credential
|
|
|
|
@staticmethod
|
|
def _get_azure_storage_access_key():
|
|
return os.getenv("AZURE_STORAGE_ACCESS_KEY")
|
|
|
|
@staticmethod
|
|
def _download_git_repo(uri: str, out_dir: str) -> str:
|
|
"""
|
|
Supports authentication via:
|
|
- Username in URL: https://username@host/repo.git
|
|
- Username from GIT_USERNAME environment variable
|
|
- Password from GIT_PASSWORD environment variable (from Kubernetes secret)
|
|
"""
|
|
from dulwich import porcelain
|
|
from dulwich.errors import GitProtocolError
|
|
from urllib.parse import urlparse, urlunparse
|
|
|
|
logger.info("Downloading Git repository %s into %s", uri, out_dir)
|
|
|
|
parsed = urlparse(uri)
|
|
username = None
|
|
clean_uri = uri
|
|
|
|
# Extract username from URL if present (format: https://username@host/repo.git)
|
|
# Note: If password is in URL (https://user:pass@host/repo.git), it will be ignored
|
|
# as passwords should come from Kubernetes secrets via GIT_PASSWORD env var
|
|
if parsed.username:
|
|
username = parsed.username
|
|
|
|
# Reconstruct URI without username/password
|
|
netloc = parsed.hostname
|
|
if parsed.port:
|
|
netloc += f":{parsed.port}"
|
|
clean_parsed = parsed._replace(netloc=netloc)
|
|
clean_uri = urlunparse(clean_parsed)
|
|
|
|
if not username:
|
|
username = os.getenv("GIT_USERNAME")
|
|
|
|
password = os.getenv("GIT_PASSWORD")
|
|
|
|
try:
|
|
clone_kwargs = {"depth": 1}
|
|
if username:
|
|
clone_kwargs["username"] = username
|
|
if password:
|
|
clone_kwargs["password"] = password
|
|
|
|
porcelain.clone(clean_uri, out_dir, **clone_kwargs)
|
|
logger.info("git clone successful")
|
|
|
|
except GitProtocolError as e:
|
|
raise_storage_error("Git", uri, e, clean_uri)
|
|
except Exception as e:
|
|
raise_storage_error("Git", uri, e, clean_uri)
|
|
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
def _download_local(
|
|
uri,
|
|
out_dir=None,
|
|
allow_patterns: Optional[List[str]] = None,
|
|
ignore_patterns: Optional[List[str]] = None,
|
|
) -> str:
|
|
local_path = uri.replace(_LOCAL_PREFIX, "", 1)
|
|
if not os.path.exists(local_path):
|
|
raise RuntimeError("Local path %s does not exist." % (uri))
|
|
|
|
if out_dir is None:
|
|
return local_path
|
|
elif not os.path.isdir(out_dir):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
if os.path.isdir(local_path):
|
|
local_path = os.path.join(local_path, "*")
|
|
|
|
file_count = 0
|
|
for src in glob.glob(local_path):
|
|
_, tail = os.path.split(src)
|
|
if not _should_download(tail, allow_patterns, ignore_patterns):
|
|
logger.info("Skipping %s due to file pattern filter", src)
|
|
continue
|
|
dest_path = os.path.join(out_dir, tail)
|
|
logger.info("Linking: %s to %s", src, dest_path)
|
|
if not os.path.exists(dest_path):
|
|
os.symlink(src, dest_path)
|
|
else:
|
|
logger.info("File %s already exist", dest_path)
|
|
file_count += 1
|
|
if file_count == 0:
|
|
raise RuntimeError("Failed to fetch model. No model found in %s." % (uri))
|
|
# Unpack compressed file, supports .tgz, tar.gz and zip file formats.
|
|
if file_count == 1:
|
|
mimetype, _ = mimetypes.guess_type(dest_path)
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
def _download_from_uri(uri, out_dir=None) -> str:
|
|
url = urlparse(uri)
|
|
filename = os.path.basename(url.path)
|
|
# Determine if the symbol '?' exists in the path
|
|
if mimetypes.guess_type(url.path)[0] is None and url.query != "":
|
|
mimetype, encoding = mimetypes.guess_type(url.query)
|
|
else:
|
|
mimetype, encoding = mimetypes.guess_type(url.path)
|
|
local_path = os.path.join(out_dir, filename)
|
|
|
|
if filename == "":
|
|
raise RuntimeError("No filename contained in URI: %s" % (uri))
|
|
|
|
# Get header information from host url
|
|
headers = {}
|
|
host_uri = url.hostname
|
|
|
|
headers_json = os.getenv(host_uri + _HEADERS_SUFFIX, "{}")
|
|
headers = json.loads(headers_json)
|
|
|
|
try:
|
|
response = requests.get(uri, stream=True, headers=headers, timeout=30)
|
|
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
|
|
raise_storage_error("HTTP", uri, e, host_uri)
|
|
|
|
with response:
|
|
check_http_response(uri, response)
|
|
zip_content_types = (
|
|
"application/x-zip-compressed",
|
|
"application/zip",
|
|
"application/zip-compressed",
|
|
)
|
|
if mimetype == "application/zip" and not response.headers.get(
|
|
"Content-Type", ""
|
|
).startswith(zip_content_types):
|
|
raise RuntimeError(
|
|
"URI: %s did not respond with any of following 'Content-Type': "
|
|
% uri
|
|
+ ", ".join(zip_content_types)
|
|
)
|
|
tar_content_types = (
|
|
"application/x-tar",
|
|
"application/x-gtar",
|
|
"application/x-gzip",
|
|
"application/gzip",
|
|
)
|
|
if mimetype == "application/x-tar" and not response.headers.get(
|
|
"Content-Type", ""
|
|
).startswith(tar_content_types):
|
|
raise RuntimeError(
|
|
"URI: %s did not respond with any of following 'Content-Type': "
|
|
% uri
|
|
+ ", ".join(tar_content_types)
|
|
)
|
|
if (
|
|
mimetype != "application/zip" and mimetype != "application/x-tar"
|
|
) and not response.headers.get("Content-Type", "").startswith(
|
|
"application/octet-stream"
|
|
):
|
|
raise RuntimeError(
|
|
"URI: %s did not respond with 'Content-Type': 'application/octet-stream'"
|
|
% uri
|
|
)
|
|
|
|
if encoding == "gzip":
|
|
stream = gzip.GzipFile(fileobj=response.raw)
|
|
local_path = os.path.join(out_dir, f"{filename}.tar")
|
|
else:
|
|
stream = response.raw
|
|
with open(local_path, "wb") as out:
|
|
shutil.copyfileobj(stream, out)
|
|
|
|
if mimetype in ["application/x-tar", "application/zip"]:
|
|
out_dir = Storage._unpack_archive_file(local_path, mimetype, out_dir)
|
|
return out_dir
|
|
|
|
@staticmethod
|
|
def _unpack_archive_file(file_path, mimetype, target_dir=None) -> str:
|
|
if not target_dir:
|
|
target_dir = os.path.dirname(file_path)
|
|
|
|
try:
|
|
logger.info("Unpacking: %s", file_path)
|
|
if mimetype == "application/x-tar":
|
|
with tarfile.open(file_path, "r", encoding="utf-8") as archive:
|
|
archive.extractall(target_dir, filter="data")
|
|
else:
|
|
with zipfile.ZipFile(file_path, "r") as archive:
|
|
archive.extractall(target_dir)
|
|
except (tarfile.TarError, zipfile.BadZipfile) as e:
|
|
raise RuntimeError(
|
|
"Failed to unpack archive file. The file format is not valid."
|
|
) from e
|
|
os.remove(file_path)
|
|
return target_dir
|