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)