def _lfs_log_progress()

in src/huggingface_hub/repository.py [0:0]
73 lines of code
16 McCabe index (conditional complexity)

def _lfs_log_progress():
    """
    This is a context manager that will log the Git LFS progress of cleaning,
    smudging, pulling and pushing.
    """

    if logger.getEffectiveLevel() >= logging.ERROR:
        try:
            yield
        except Exception:
            pass
        return

    def output_progress(stopping_event: threading.Event):
        """
        To be launched as a separate thread with an event meaning it should stop
        the tail.
        """
        # Key is tuple(state, filename), value is a dict(tqdm bar and a previous value)
        pbars: Dict[Tuple[str, str], PbarT] = {}

        def close_pbars():
            for pbar in pbars.values():
                pbar["bar"].update(pbar["bar"].total - pbar["past_bytes"])
                pbar["bar"].refresh()
                pbar["bar"].close()

        def tail_file(filename) -> Iterator[str]:
            """
            Creates a generator to be iterated through, which will return each
            line one by one. Will stop tailing the file if the stopping_event is
            set.
            """
            with open(filename, "r") as file:
                current_line = ""
                while True:
                    if stopping_event.is_set():
                        close_pbars()
                        break

                    line_bit = file.readline()
                    if line_bit is not None and not len(line_bit.strip()) == 0:
                        current_line += line_bit
                        if current_line.endswith("\n"):
                            yield current_line
                            current_line = ""
                    else:
                        time.sleep(1)

        # If the file isn't created yet, wait for a few seconds before trying again.
        # Can be interrupted with the stopping_event.
        while not os.path.exists(os.environ["GIT_LFS_PROGRESS"]):
            if stopping_event.is_set():
                close_pbars()
                return

            time.sleep(2)

        for line in tail_file(os.environ["GIT_LFS_PROGRESS"]):
            try:
                state, file_progress, byte_progress, filename = line.split()
            except ValueError as error:
                # Try/except to ease debugging. See https://github.com/huggingface/huggingface_hub/issues/1373.
                raise ValueError(f"Cannot unpack LFS progress line:\n{line}") from error
            description = f"{state.capitalize()} file {filename}"

            current_bytes, total_bytes = byte_progress.split("/")
            current_bytes_int = int(current_bytes)
            total_bytes_int = int(total_bytes)

            pbar = pbars.get((state, filename))
            if pbar is None:
                # Initialize progress bar
                pbars[(state, filename)] = {
                    "bar": tqdm(
                        desc=description,
                        initial=current_bytes_int,
                        total=total_bytes_int,
                        unit="B",
                        unit_scale=True,
                        unit_divisor=1024,
                        name="huggingface_hub.lfs_upload",
                    ),
                    "past_bytes": int(current_bytes),
                }
            else:
                # Update progress bar
                pbar["bar"].update(current_bytes_int - pbar["past_bytes"])
                pbar["past_bytes"] = current_bytes_int

    current_lfs_progress_value = os.environ.get("GIT_LFS_PROGRESS", "")

    with SoftTemporaryDirectory() as tmpdir:
        os.environ["GIT_LFS_PROGRESS"] = os.path.join(tmpdir, "lfs_progress")
        logger.debug(f"Following progress in {os.environ['GIT_LFS_PROGRESS']}")

        exit_event = threading.Event()
        x = threading.Thread(target=output_progress, args=(exit_event,), daemon=True)
        x.start()

        try:
            yield
        finally:
            exit_event.set()
            x.join()

            os.environ["GIT_LFS_PROGRESS"] = current_lfs_progress_value