implementation of a custom transfer agent for the transfer type multipart for gitlfs inspired by github comcbartzgitlfsswifttransferagentblobmastergitlfsswifttransfer py spec is github comgitlfsgitlfsblobmasterdocscustomtransfers md to launch debugger while developing lfs customtransfer multipart path pathtotransformers envbinpython args m debugpy listen 5678 waitforclient pathtotransformerssrctransformerscommandstransformerscli py lfsmultipartupload import json import os import subprocess import sys import warnings from argparse import argumentparser from contextlib import abstractcontextmanager from typing import dict list optional import requests from utils import logging from import basetransformersclicommand logger logging getloggername pylint disableinvalidname lfsmultipartuploadcommand lfsmultipartupload class lfscommandsbasetransformersclicommand staticmethod def registersubcommandparser argumentparser enableparser parser addparser lfsenablelargefiles help deprecated use huggingfacecli instead configure your repository to enable upload of files 5gb enableparser addargumentpath typestr helplocal path to repository you want to configure enableparser setdefaultsfunclambda args lfsenablecommandargs uploadparser parser addparser lfsmultipartuploadcommand help deprecated use huggingfacecli instead command will get called by gitlfs do not call it directly uploadparser setdefaultsfunclambda args lfsuploadcommandargs class lfsenablecommand def initself args self args args def runself warnings warn managing repositories through transformerscli is deprecated please use huggingfacecli instead localpath os path abspathself args path if not os path isdirlocalpath printthis does not look like a valid git repo exit1 subprocess run git config lfs customtransfer multipart path transformerscli split checktrue cwdlocalpath subprocess run fgit config lfs customtransfer multipart args lfsmultipartuploadcommand split checktrue cwdlocalpath printlocal repo set up for largefiles def writemsgmsg dict read line delimited json from stdin msg json loadssys stdin readline strip if terminate in msg gettype msg getevent terminate message received return none if msg getevent not in download upload logger criticalreceived unexpected message sys exit1 return msg class filesliceabstractcontextmanager def initself filepath str seekfrom int readlimit int self filepath filepath self seekfrom seekfrom self readlimit readlimit self nseen 0 def enterself self f openself filepath rb self f seekself seekfrom return self def lenself totallength os fstatself f fileno stsize return minself readlimit totallength self seekfrom def readself n1 if self nseen self readlimit return b remainingamount self readlimit self nseen data self f readremainingamount if n 0 else minn remainingamount self nseen lendata return data def iterself yield self readn4 1024 1024 def exitself args self f close class lfsuploadcommand def initself args self args args def runself immediately after invoking a custom transfer process gitlfs sends initiation data to the process over stdin this tells the process useful information about the configuration initmsg json loadssys stdin readline strip if not initmsg getevent init and initmsg getoperation upload writemsgerror code 32 message wrong lfs init operation sys exit1 the transfer process should use the information it needs from the initiation structure and also perform any oneoff setup tasks it needs to do it should then respond on stdout with a simple empty confirmation structure as follows writemsg after the initiation exchange gitlfs will send any number of transfer requests to the stdin of the transfer process in a serial sequence while true msg readmsg if msg is none when all transfers have been processed gitlfs will send a terminate event to the stdin of the transfer process on receiving this message the transfer process should clean up and terminate no response is expected sys exit0 oid msgoid filepath msgpath completionurl msgactionhref header msgactionheader chunksize intheader popchunksize presignedurls liststr listheader values parts for i presignedurl in enumeratepresignedurls with fileslicefilepath seekfromi chunksize readlimitchunksize as data r requests putpresignedurl datadata r raiseforstatus parts append etag r headers getetag partnumber i 1 in order to support progress reporting while data is uploading downloading the transfer process should post messages to stdout writemsg event progress oid oid bytessofar i 1 chunksize bytessincelast chunksize not precise but that s ok r requests post completionurl json oid oid parts parts r raiseforstatus writemsgevent complete oid oid pylint disable invalid name implementation of a custom transfer agent for the transfer type multipart for git lfs this lets users upload large files 5gb spec for lfs custom transfer agent is https github com git lfs git lfs blob master docs custom transfers md this introduces two commands to the cli 1 transformers cli lfs enable largefiles this should be executed once for each model repo that contains a model file 5gb it s documented in the error message you get if you just try to git push a 5gb file without having enabled it before 2 transformers cli lfs multipart upload this command is called by lfs directly and is not meant to be called by the user write out the message in line delimited json read line delimited json from stdin terminate message received file like object that only reads a slice of a file inspired by stackoverflow com a 29838711 593036 immediately after invoking a custom transfer process git lfs sends initiation data to the process over stdin this tells the process useful information about the configuration the transfer process should use the information it needs from the initiation structure and also perform any one off setup tasks it needs to do it should then respond on stdout with a simple empty confirmation structure as follows after the initiation exchange git lfs will send any number of transfer requests to the stdin of the transfer process in a serial sequence when all transfers have been processed git lfs will send a terminate event to the stdin of the transfer process on receiving this message the transfer process should clean up and terminate no response is expected in order to support progress reporting while data is uploading downloading the transfer process should post messages to stdout not precise but that s ok
import json
import os
import subprocess
import sys
import warnings
from argparse import ArgumentParser
from contextlib import AbstractContextManager
from typing import Dict, List, Optional

import requests

from ..utils import logging
from . import BaseTransformersCLICommand


logger = logging.get_logger(__name__)  


LFS_MULTIPART_UPLOAD_COMMAND = "lfs-multipart-upload"


class LfsCommands(BaseTransformersCLICommand):
    

    @staticmethod
    def register_subcommand(parser: ArgumentParser):
        enable_parser = parser.add_parser(
            "lfs-enable-largefiles",
            help=(
                "Deprecated: use `huggingface-cli` instead. Configure your repository to enable upload of files > 5GB."
            ),
        )
        enable_parser.add_argument("path", type=str, help="Local path to repository you want to configure.")
        enable_parser.set_defaults(func=lambda args: LfsEnableCommand(args))

        upload_parser = parser.add_parser(
            LFS_MULTIPART_UPLOAD_COMMAND,
            help=(
                "Deprecated: use `huggingface-cli` instead. "
                "Command will get called by git-lfs, do not call it directly."
            ),
        )
        upload_parser.set_defaults(func=lambda args: LfsUploadCommand(args))


class LfsEnableCommand:
    def __init__(self, args):
        self.args = args

    def run(self):
        warnings.warn(
            "Managing repositories through transformers-cli is deprecated. Please use `huggingface-cli` instead."
        )
        local_path = os.path.abspath(self.args.path)
        if not os.path.isdir(local_path):
            print("This does not look like a valid git repo.")
            exit(1)
        subprocess.run(
            "git config lfs.customtransfer.multipart.path transformers-cli".split(), check=True, cwd=local_path
        )
        subprocess.run(
            f"git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}".split(),
            check=True,
            cwd=local_path,
        )
        print("Local repo set up for largefiles")


def write_msg(msg: Dict):
    
    msg = json.dumps(msg) + "\n"
    sys.stdout.write(msg)
    sys.stdout.flush()


def read_msg() -> Optional[Dict]:
    
    msg = json.loads(sys.stdin.readline().strip())

    if "terminate" in (msg.get("type"), msg.get("event")):
        
        return None

    if msg.get("event") not in ("download", "upload"):
        logger.critical("Received unexpected message")
        sys.exit(1)

    return msg


class FileSlice(AbstractContextManager):
    

    def __init__(self, filepath: str, seek_from: int, read_limit: int):
        self.filepath = filepath
        self.seek_from = seek_from
        self.read_limit = read_limit
        self.n_seen = 0

    def __enter__(self):
        self.f = open(self.filepath, "rb")
        self.f.seek(self.seek_from)
        return self

    def __len__(self):
        total_length = os.fstat(self.f.fileno()).st_size
        return min(self.read_limit, total_length - self.seek_from)

    def read(self, n=-1):
        if self.n_seen >= self.read_limit:
            return b""
        remaining_amount = self.read_limit - self.n_seen
        data = self.f.read(remaining_amount if n < 0 else min(n, remaining_amount))
        self.n_seen += len(data)
        return data

    def __iter__(self):
        yield self.read(n=4 * 1024 * 1024)

    def __exit__(self, *args):
        self.f.close()


class LfsUploadCommand:
    def __init__(self, args):
        self.args = args

    def run(self):
        
        
        
        init_msg = json.loads(sys.stdin.readline().strip())
        if not (init_msg.get("event") == "init" and init_msg.get("operation") == "upload"):
            write_msg({"error": {"code": 32, "message": "Wrong lfs init operation"}})
            sys.exit(1)

        
        
        
        
        write_msg({})

        
        
        while True:
            msg = read_msg()
            if msg is None:
                
                
                
                
                sys.exit(0)

            oid = msg["oid"]
            filepath = msg["path"]
            completion_url = msg["action"]["href"]
            header = msg["action"]["header"]
            chunk_size = int(header.pop("chunk_size"))
            presigned_urls: List[str] = list(header.values())

            parts = []
            for i, presigned_url in enumerate(presigned_urls):
                with FileSlice(filepath, seek_from=i * chunk_size, read_limit=chunk_size) as data:
                    r = requests.put(presigned_url, data=data)
                    r.raise_for_status()
                    parts.append(
                        {
                            "etag": r.headers.get("etag"),
                            "partNumber": i + 1,
                        }
                    )
                    
                    
                    write_msg(
                        {
                            "event": "progress",
                            "oid": oid,
                            "bytesSoFar": (i + 1) * chunk_size,
                            "bytesSinceLast": chunk_size,
                        }
                    )
                    

            r = requests.post(
                completion_url,
                json={
                    "oid": oid,
                    "parts": parts,
                },
            )
            r.raise_for_status()

            write_msg({"event": "complete", "oid": oid})