MLX distributed file transfer

README.md

Transfer files with MLX distributed

This script enables transfering files from one Mac to another using MLX distributed.

First configure the nodes if you haven't already

mlx.distributed_config --verbose --hosts hostname1,hostname2 --backend jaccl \
    --env MLX_METAL_FAST_SYNCH=1 --auto-setup --output hostfile.json

then all you have to from one of the two nodes

mlx.launch --verbose --hostfile hostfile.json -- file_transfer.py <src-path> <dst-path>

You can also initiate the copy from another node

mlx.launch --verbose --hostfile hostfile.json --no-verify -- \
    /full/path/to/python file_transfer.py <src-path> <dst-path>

Example

In the following I am transferring a 5GB file from one node to another

$ mlx.launch --verbose --hostfile m3-ultra-2x.json --no-verify -- \
    /Users/angelos/miniconda3/envs/mlx/bin/python file_transfer.py --from 1 --to 0 /tmp/src.gz /tmp/dst.gz
[INFO] Running /Users/angelos/miniconda3/envs/mlx/bin/python file_transfer.py --from 1 --to 0 /tmp/src.gz /tmp/dst.gz
Sending /tmp/src.gz:  87% 4.20G/4.83G [00:00<00:00, 6.49GB/s]Received 4826186110 bytes in 0.77s (5942.16 MB/s)
Sending /tmp/src.gz: 100% 4.83G/4.83G [00:00<00:00, 6.19GB/s]
Sent 4826186110 bytes in 0.78s (5900.61 MB/s)
[INFO] Node with rank 1 completed
[INFO] Node with rank 0 completed

file_transfer.py

import argparse
import time

import mlx.core as mx
from tqdm import tqdm


def transfer_file(from_rank, to_rank, input_file, output_file, chunk_size):
    world = mx.distributed.init()

    if world.rank() == from_rank:
        with open(input_file, "rb") as f:
            f.seek(0, 2)
            total_size = f.tell()
            f.seek(0)

            pbar = tqdm(
                total=total_size,
                unit="B",
                unit_scale=True,
                desc=f"Sending {input_file}",
            )
            start_time = time.time()

            while True:
                data = f.read(chunk_size)
                if not data:
                    mx.eval(mx.distributed.send(0, to_rank))
                    break

                mx.eval(mx.distributed.send(len(data), to_rank))
                mx.async_eval(mx.distributed.send(mx.array(data), to_rank))
                pbar.update(len(data))

            elapsed = time.time() - start_time
            speed = total_size / elapsed if elapsed > 0 else 0
            pbar.close()
            print(
                f"Sent {total_size} bytes in {elapsed:.2f}s "
                f"({speed/1024/1024:.2f} MB/s)"
            )

    elif world.rank() == to_rank:
        with open(output_file, "wb") as f:
            total_size = 0
            start_time = time.time()

            chunk_size = mx.distributed.recv((1,), mx.int32, from_rank).item()
            total_size += chunk_size
            data = None
            if chunk_size > 0:
                data = mx.distributed.recv((chunk_size,), mx.uint8, from_rank)
                mx.eval(data)

            while chunk_size > 0:
                chunk_size = mx.distributed.recv((1,), mx.int32, from_rank).item()
                total_size += chunk_size
                next_data = None
                if chunk_size > 0:
                    next_data = mx.distributed.recv((chunk_size,), mx.uint8, from_rank)
                    mx.async_eval(next_data)

                f.write(bytes(data))
                data = next_data

            elapsed = time.time() - start_time
            speed = total_size / elapsed if elapsed > 0 else 0
            print(
                f"Received {total_size} bytes in {elapsed:.2f}s "
                f"({speed/1024/1024:.2f} MB/s)"
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Transfer a file between ranks using MLX distributed"
    )
    parser.add_argument("source", help="Source file path")
    parser.add_argument("destination", nargs="?", help="Destination file path")
    parser.add_argument(
        "--from", type=int, default=0, dest="from_rank", help="Source rank (default: 0)"
    )
    parser.add_argument(
        "--to", type=int, default=1, dest="to_rank", help="Destination rank (default: 1)"
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=100 * 1024 * 1024,
        help="Chunk size in bytes (default: 100MB)",
    )
    args = parser.parse_args()

    output_file = args.destination if args.destination else args.source
    transfer_file(
        from_rank=args.from_rank,
        to_rank=args.to_rank,
        input_file=args.source,
        output_file=output_file,
        chunk_size=args.chunk_size,
    )