Ejemplo n.º 1
0
async def handle_posted_data(
    content_file: UploadFile,
    root_path: str,
    path: str,
    upload: bool,
    is_file: bool,
    overwrite: bool = True,
    untar: bool = True,
) -> str:
    tmp_path = "{}/{}".format(root_path, os.path.basename(
        content_file.filename)).rstrip("/")
    if path:
        root_path = "{}/{}".format(root_path, path).rstrip("/")
        if is_file:
            root_path = "{}/{}".format(root_path,
                                       os.path.basename(content_file.filename))
    else:
        if untar:
            root_path = "{}/{}".format(root_path,
                                       DEFAULT_UPLOADS_PATH).rstrip("/")
        else:
            root_path = tmp_path
    if not untar:
        tmp_path = root_path
    full_tmppath = os.path.join(settings.AGENT_CONFIG.artifacts_root, tmp_path)
    full_filepath = os.path.join(settings.AGENT_CONFIG.artifacts_root,
                                 root_path)

    if overwrite and os.path.exists(full_filepath):
        delete_path(full_filepath)
    if not overwrite and os.path.exists(full_filepath):
        return full_filepath
    # Always clean tmp path
    if overwrite and os.path.exists(full_tmppath):
        delete_path(full_tmppath)

    check_or_create_path(full_tmppath, is_dir=False)
    check_or_create_path(full_filepath, is_dir=not is_file)

    # Creating the new file
    with open(full_tmppath, "wb") as destination:
        for chunk in content_file.file:
            destination.write(chunk)
    if untar:
        untar_file(full_tmppath,
                   extract_path=full_filepath,
                   use_filepath=False)
    if upload:
        if is_file:
            await upload_file(subpath=root_path)
        else:
            await upload_dir(subpath=root_path)
    return root_path
Ejemplo n.º 2
0
    def download(
        self,
        url,
        filename,
        params=None,
        headers=None,
        timeout=None,
        session=None,
        untar=False,
        delete_tar=True,
        extract_path=None,
    ):
        """
        Download the file from the given url at the current path
        """
        # pylint:disable=too-many-branches
        logger.debug("Downloading files from url: %s", url)

        request_headers = self._get_headers(headers=headers)
        timeout = timeout if timeout is not None else settings.LONG_REQUEST_TIMEOUT
        session = session or requests.Session()

        try:
            response = session.get(
                url=url,
                params=params,
                headers=request_headers,
                timeout=timeout,
                stream=True,
            )

            self.check_response_status(response, url)
            with open(filename, "wb") as f:
                # chunk mode response doesn't have content-length so we are
                # using a custom header here
                content_length = response.headers.get("x-polyaxon-content-length")
                if not content_length:
                    content_length = response.headers.get("content-length")
                if content_length:
                    for chunk in progress_bar(
                        response.iter_content(chunk_size=1024),
                        expected_size=(int(content_length) / 1024) + 1,
                    ):
                        if chunk:
                            f.write(chunk)
                else:
                    for chunk in response.iter_content(chunk_size=1024):
                        if chunk:
                            f.write(chunk)

            if untar:
                filename = untar_file(
                    filename=filename, delete_tar=delete_tar, extract_path=extract_path
                )
            return filename
        except (
            requests.exceptions.ConnectionError,
            requests.exceptions.RequestException,
            requests.exceptions.Timeout,
            requests.exceptions.HTTPError,
        ) as exception:
            try:
                logger.debug("Exception: %s", exception)
            except TypeError:
                pass

            raise PolyaxonShouldExitError(
                "Error connecting to Polyaxon server on `{}`.\n"
                "An Error `{}` occurred.\n"
                "Check your host and ports configuration "
                "and your internet connection.".format(url, exception)
            )