Function download_model_from_s3

python/ray/llm/tests/conftest.py:19–41

Download the model checkpoint and tokenizer from S3 for testing The reason to download the model from S3 is to avoid downloading the model from HuggingFace hub during testing, which is flaky because of the rate limit and HF hub downtime. Args: remote_url: The remote URL

(
    remote_url: str, file_list: List[str]
)

def download_model_from_s3(
    remote_url: str, file_list: List[str]
) -> Generator[str, None, None]:
    """
    Download the model checkpoint and tokenizer from S3 for testing
    The reason to download the model from S3 is to avoid downloading the model
    from HuggingFace hub during testing, which is flaky because of the rate
    limit and HF hub downtime.

    Args:
        remote_url: The remote URL to download the model from.
        file_list: The list of files to download.

    Yields:
        str: The path to the downloaded model checkpoint and tokenizer.
    """
    with tempfile.TemporaryDirectory(prefix="ray-llm-test-model") as checkpoint_dir:
        print(f"Downloading model from {remote_url} to {checkpoint_dir}", flush=True)
        for file_name in file_list:
            response = requests.get(remote_url + file_name)
            with open(os.path.join(checkpoint_dir, file_name), "wb") as fp:
                fp.write(response.content)
        yield os.path.abspath(checkpoint_dir)


@pytest.fixture(scope="session")