def put_block_list(sas_url, block_list, headers):
"""
Performs an Azure `Put Block List` operation
(https://docs.microsoft.com/en-us/rest/api/storageservices/put-block-list)
:param sas_url: A shared access signature URL referring to the Azure Block Blob
to which the specified data should be staged.
:param block_list: A list of uncommitted base64-encoded string block IDs to commit. For
more information, see
https://docs.microsoft.com/en-us/rest/api/storageservices/put-block-list.
:param headers: Headers to include in the Put Block request body.
"""
request_url = _append_query_parameters(sas_url, {"comp": "blocklist"})
data = _build_block_list_xml(block_list)
request_headers = {}
for name, value in headers.items():
if _is_valid_put_block_list_header(name):
request_headers[name] = value
else:
_logger.debug(
"Removed unsupported '%s' header for Put Block List operation",
name)
with rest_utils.cloud_storage_http_request(
"put", request_url, data=data,
headers=request_headers) as response:
rest_utils.augmented_raise_for_status(response)
def put_block(sas_url, block_id, data, headers):
"""
Performs an Azure `Put Block` operation
(https://docs.microsoft.com/en-us/rest/api/storageservices/put-block)
:param sas_url: A shared access signature URL referring to the Azure Block Blob
to which the specified data should be staged.
:param block_id: A base64-encoded string identifying the block.
:param data: Data to include in the Put Block request body.
:param headers: Additional headers to include in the Put Block request body
(the `x-ms-blob-type` header is always included automatically).
"""
request_url = _append_query_parameters(sas_url, {
"comp": "block",
"blockid": block_id
})
request_headers = deepcopy(_PUT_BLOCK_HEADERS)
for name, value in headers.items():
if _is_valid_put_block_header(name):
request_headers[name] = value
else:
_logger.debug(
"Removed unsupported '%s' header for Put Block operation",
name)
with rest_utils.cloud_storage_http_request(
"put", request_url, data=data,
headers=request_headers) as response:
rest_utils.augmented_raise_for_status(response)
def _download_file(self, remote_file_path, local_path):
url = posixpath.join(self.artifact_uri, remote_file_path)
with self._session.get(url, stream=True, timeout=10) as resp:
augmented_raise_for_status(resp)
with open(local_path, "wb") as f:
chunk_size = 1024 * 1024 # 1 MB
for chunk in resp.iter_content(chunk_size=chunk_size):
f.write(chunk)
def log_artifact(self, local_file, artifact_path=None):
verify_artifact_path(artifact_path)
file_name = os.path.basename(local_file)
paths = (artifact_path, file_name) if artifact_path else (file_name, )
url = posixpath.join(self.artifact_uri, *paths)
with open(local_file, "rb") as f:
resp = self._session.put(url, data=f, timeout=600)
augmented_raise_for_status(resp)
def test_mlflow_server_command(command):
port = get_safe_port()
cmd = ["mlflow", command, "--port", str(port)]
process = subprocess.Popen(cmd)
try:
_await_server_up_or_die(port, timeout=10)
resp = requests.get(f"http://localhost:{port}/health")
augmented_raise_for_status(resp)
assert resp.text == "OK"
finally:
process.kill()
def _download_file(self, remote_file_path, local_path):
endpoint = posixpath.join("/", remote_file_path)
resp = http_request(self._host_creds,
endpoint,
"GET",
stream=True,
timeout=10)
augmented_raise_for_status(resp)
with open(local_path, "wb") as f:
chunk_size = 1024 * 1024 # 1 MB
for chunk in resp.iter_content(chunk_size=chunk_size):
f.write(chunk)
def test_mlflow_artifact_list_in_artifacts_only_mode():
port = get_safe_port()
cmd = ["mlflow", "server", "--port", str(port), "--artifacts-only", "--serve-artifacts"]
process = subprocess.Popen(cmd)
try:
_await_server_up_or_die(port, timeout=10)
resp = requests.get(f"http://localhost:{port}/api/2.0/mlflow-artifacts/artifacts")
augmented_raise_for_status(resp)
assert resp.status_code == 200
assert resp.text == "{}"
finally:
process.kill()
def log_artifact(self, local_file, artifact_path=None):
verify_artifact_path(artifact_path)
file_name = os.path.basename(local_file)
paths = (artifact_path, file_name) if artifact_path else (file_name, )
endpoint = posixpath.join("/", *paths)
with open(local_file, "rb") as f:
resp = http_request(self._host_creds,
endpoint,
"PUT",
data=f,
timeout=600)
augmented_raise_for_status(resp)
def download_file_using_http_uri(http_uri, download_path, chunk_size=100000000):
"""
Downloads a file specified using the `http_uri` to a local `download_path`. This function
uses a `chunk_size` to ensure an OOM error is not raised a large file is downloaded.
Note : This function is meant to download files using presigned urls from various cloud
providers.
"""
with cloud_storage_http_request("get", http_uri, stream=True) as response:
augmented_raise_for_status(response)
with open(download_path, "wb") as output_file:
for chunk in response.iter_content(chunk_size=chunk_size):
if not chunk:
break
output_file.write(chunk)
def _fetch_zip_repo(uri):
import requests
from io import BytesIO
# TODO (dbczumar): Replace HTTP resolution via ``requests.get`` with an invocation of
# ```mlflow.data.download_uri()`` when the API supports the same set of available stores as
# the artifact repository (Azure, FTP, etc). See the following issue:
# https://github.com/mlflow/mlflow/issues/763.
response = requests.get(uri)
try:
augmented_raise_for_status(response)
except requests.HTTPError as error:
raise ExecutionException("Unable to retrieve ZIP file. Reason: %s" %
str(error))
return BytesIO(response.content)
def list_artifacts(self, path=None):
sep = "/mlflow-artifacts/artifacts"
head, tail = self.artifact_uri.split(sep, maxsplit=1)
url = head + sep
root = tail.lstrip("/")
params = {"path": posixpath.join(root, path) if path else root}
resp = self._session.get(url, params=params, timeout=10)
augmented_raise_for_status(resp)
file_infos = []
for f in resp.json().get("files", []):
file_info = FileInfo(
posixpath.join(path, f["path"]) if path else f["path"],
f["is_dir"],
int(f["file_size"]) if ("file_size" in f) else None,
)
file_infos.append(file_info)
return sorted(file_infos, key=lambda f: f.path)
def _signed_url_upload_file(self, credentials, local_file):
try:
headers = self._extract_headers_from_credentials(
credentials.headers)
signed_write_uri = credentials.signed_uri
# Putting an empty file in a request by reading file bytes gives 501 error.
if os.stat(local_file).st_size == 0:
with rest_utils.cloud_storage_http_request(
"put", signed_write_uri, data="",
headers=headers) as response:
augmented_raise_for_status(response)
else:
with open(local_file, "rb") as file:
with rest_utils.cloud_storage_http_request(
"put", signed_write_uri, data=file,
headers=headers) as response:
augmented_raise_for_status(response)
except Exception as err:
raise MlflowException(err)
def list_artifacts(self, path=None):
endpoint = "/mlflow-artifacts/artifacts"
url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
root = tail.lstrip("/")
params = {"path": posixpath.join(root, path) if path else root}
host_creds = _get_default_host_creds(url)
resp = http_request(host_creds,
endpoint,
"GET",
params=params,
timeout=10)
augmented_raise_for_status(resp)
file_infos = []
for f in resp.json().get("files", []):
file_info = FileInfo(
posixpath.join(path, f["path"]) if path else f["path"],
f["is_dir"],
int(f["file_size"]) if ("file_size" in f) else None,
)
file_infos.append(file_info)
return sorted(file_infos, key=lambda f: f.path)