Merge pull request #1847 from transformerlab/add/chunk-file-uploads

Chunked file uploads

Author: deep1401

api/api.py

@@ -50,6 +50,7 @@
     asset_versions,
     trackio,
     permissions,
+    upload,
 )
 from transformerlab.routers.auth import get_user_and_team  # noqa: E402
 
@@ -147,6 +148,12 @@ async def lifespan(app: FastAPI):
         )
 
         await start_remote_job_queue_worker()
+
+        # Sweep abandoned chunked-upload staging dirs older than 24 h
+        from transformerlab.services.upload_service import sweep_expired_uploads
+
+        swept = await asyncio.to_thread(sweep_expired_uploads)
+        print(f"✅ Upload staging sweep: removed {swept} expired upload(s)")
     print("FastAPI LIFESPAN: 🏁 🏁 🏁 Begin API Server 🏁 🏁 🏁", flush=True)
     yield
     # Do the following at API Shutdown:
@@ -315,6 +322,7 @@ async def validation_exception_handler(request, exc):
 app.include_router(asset_versions.router, dependencies=[Depends(get_user_and_team)])
 app.include_router(trackio.router, dependencies=[Depends(get_user_and_team)])
 app.include_router(permissions.router, dependencies=[Depends(get_user_and_team)])
+app.include_router(upload.router, dependencies=[Depends(get_user_and_team)])
 
 
 # @app.get("/")

api/test/api/test_upload.py

@@ -0,0 +1,77 @@
+import os
+
+
+def test_init_upload(client):
+    resp = client.post("/upload/init", json={"filename": "model.zip", "total_size": 1024})
+    assert resp.status_code == 200
+    body = resp.json()
+    assert "upload_id" in body
+    assert body["chunk_size"] == 64 * 1024 * 1024
+
+
+def test_chunk_upload_and_status(client):
+    init = client.post("/upload/init", json={"filename": "data.bin", "total_size": 6}).json()
+    uid = init["upload_id"]
+
+    resp = client.put(
+        f"/upload/{uid}/chunk?chunk_index=0",
+        content=b"ABC",
+        headers={"Content-Type": "application/octet-stream"},
+    )
+    assert resp.status_code == 200
+    assert resp.json()["received"] == [0]
+
+    status = client.get(f"/upload/{uid}/status").json()
+    assert status["received"] == [0]
+    assert status["complete"] is False
+
+
+def test_complete_upload(client):
+    init = client.post("/upload/init", json={"filename": "data.bin", "total_size": 6}).json()
+    uid = init["upload_id"]
+
+    client.put(
+        f"/upload/{uid}/chunk?chunk_index=0",
+        content=b"ABC",
+        headers={"Content-Type": "application/octet-stream"},
+    )
+    client.put(
+        f"/upload/{uid}/chunk?chunk_index=1",
+        content=b"DEF",
+        headers={"Content-Type": "application/octet-stream"},
+    )
+
+    resp = client.post(f"/upload/{uid}/complete", json={"total_chunks": 2})
+    assert resp.status_code == 200
+    body = resp.json()
+    assert "temp_path" in body
+
+    assert os.path.isfile(body["temp_path"])
+    with open(body["temp_path"], "rb") as f:
+        assert f.read() == b"ABCDEF"
+
+
+def test_complete_fails_with_missing_chunks(client):
+    init = client.post("/upload/init", json={"filename": "data.bin", "total_size": 6}).json()
+    uid = init["upload_id"]
+    client.put(
+        f"/upload/{uid}/chunk?chunk_index=0",
+        content=b"ABC",
+        headers={"Content-Type": "application/octet-stream"},
+    )
+
+    resp = client.post(f"/upload/{uid}/complete", json={"total_chunks": 2})
+    assert resp.status_code == 400
+    assert "Missing chunks" in resp.json()["detail"]
+
+
+def test_delete_upload(client):
+    init = client.post("/upload/init", json={"filename": "data.bin", "total_size": 1}).json()
+    uid = init["upload_id"]
+
+    resp = client.delete(f"/upload/{uid}")
+    assert resp.status_code == 200
+    assert resp.json()["status"] == "deleted"
+
+    status = client.get(f"/upload/{uid}/status")
+    assert status.status_code == 404

api/test/api/test_upload_service.py

@@ -0,0 +1,196 @@
+import json
+import os
+import pytest
+
+
+# Point staging at a temp dir for tests
+@pytest.fixture(autouse=True)
+def isolated_staging(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        "transformerlab.services.upload_service.STAGING_ROOT",
+        str(tmp_path / "staging"),
+    )
+    yield
+
+
+def test_init_creates_staging_dir():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, STAGING_ROOT
+
+    result = asyncio.run(init_upload("model.zip", 1024 * 1024))
+
+    assert "upload_id" in result
+    assert result["chunk_size"] == 64 * 1024 * 1024
+    staging = os.path.join(STAGING_ROOT, result["upload_id"])
+    assert os.path.isdir(staging)
+    with open(os.path.join(staging, "meta.json")) as f:
+        meta = json.load(f)
+    assert meta["filename"] == "model.zip"
+    assert meta["total_size"] == 1024 * 1024
+
+
+def test_save_chunk_writes_file():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk, STAGING_ROOT
+
+    r = asyncio.run(init_upload("data.bin", 200))
+    uid = r["upload_id"]
+
+    received = asyncio.run(save_chunk(uid, 0, b"hello chunk 0"))
+
+    chunk_path = os.path.join(STAGING_ROOT, uid, "0")
+    assert os.path.isfile(chunk_path)
+    assert open(chunk_path, "rb").read() == b"hello chunk 0"
+    assert received == [0]
+
+
+def test_save_chunk_idempotent():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk
+
+    r = asyncio.run(init_upload("data.bin", 200))
+    uid = r["upload_id"]
+
+    asyncio.run(save_chunk(uid, 0, b"first write"))
+    asyncio.run(save_chunk(uid, 0, b"second write"))
+
+    from transformerlab.services.upload_service import STAGING_ROOT
+
+    chunk_path = os.path.join(STAGING_ROOT, uid, "0")
+    assert open(chunk_path, "rb").read() == b"second write"
+
+
+def test_get_status_returns_received_chunks():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk, get_status
+
+    r = asyncio.run(init_upload("data.bin", 300))
+    uid = r["upload_id"]
+    asyncio.run(save_chunk(uid, 0, b"a"))
+    asyncio.run(save_chunk(uid, 2, b"c"))
+
+    status = asyncio.run(get_status(uid))
+
+    assert status["upload_id"] == uid
+    assert status["received"] == [0, 2]
+    assert status["complete"] is False
+
+
+def test_assemble_upload_concatenates_chunks():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk, assemble_upload_sync, STAGING_ROOT
+
+    r = asyncio.run(init_upload("data.bin", 6))
+    uid = r["upload_id"]
+    asyncio.run(save_chunk(uid, 0, b"ABC"))
+    asyncio.run(save_chunk(uid, 1, b"DEF"))
+
+    path = assemble_upload_sync(uid, total_chunks=2)
+
+    assert open(path, "rb").read() == b"ABCDEF"
+    # _staging_dir returns os.path.realpath(), so compare against the resolved path.
+    assert path == os.path.realpath(os.path.join(STAGING_ROOT, uid, "assembled"))
+
+
+def test_assemble_upload_raises_on_missing_chunks():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk, assemble_upload_sync
+
+    r = asyncio.run(init_upload("data.bin", 9))
+    uid = r["upload_id"]
+    asyncio.run(save_chunk(uid, 0, b"AAA"))
+    # chunk 1 is missing
+
+    with pytest.raises(ValueError, match="Missing chunks"):
+        assemble_upload_sync(uid, total_chunks=2)
+
+
+def test_delete_upload_removes_dir():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, delete_upload, STAGING_ROOT
+
+    r = asyncio.run(init_upload("data.bin", 10))
+    uid = r["upload_id"]
+
+    asyncio.run(delete_upload(uid))
+
+    assert not os.path.isdir(os.path.join(STAGING_ROOT, uid))
+
+
+def test_get_assembled_path_returns_path_when_assembled():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, save_chunk, assemble_upload_sync, get_assembled_path
+
+    r = asyncio.run(init_upload("data.bin", 3))
+    uid = r["upload_id"]
+    asyncio.run(save_chunk(uid, 0, b"XYZ"))
+    assembled = assemble_upload_sync(uid, total_chunks=1)
+
+    result = asyncio.run(get_assembled_path(uid))
+
+    assert result == assembled
+
+
+def test_get_assembled_path_raises_when_not_assembled():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, get_assembled_path
+
+    r = asyncio.run(init_upload("data.bin", 3))
+    uid = r["upload_id"]
+
+    with pytest.raises(ValueError, match="has not been assembled"):
+        asyncio.run(get_assembled_path(uid))
+
+
+def test_save_chunk_raises_on_unknown_upload_id():
+    import asyncio
+    from transformerlab.services.upload_service import save_chunk
+
+    # Must be a valid 32-char hex UUID format but refer to an upload that doesn't exist.
+    with pytest.raises(ValueError, match="not found"):
+        asyncio.run(save_chunk("a" * 32, 0, b"data"))
+
+
+def test_save_chunk_raises_on_invalid_upload_id():
+    import asyncio
+    from transformerlab.services.upload_service import save_chunk
+
+    with pytest.raises(ValueError, match="Invalid upload_id"):
+        asyncio.run(save_chunk("../../etc/passwd", 0, b"data"))
+
+
+def test_get_filename_returns_filename():
+    import asyncio
+    from transformerlab.services.upload_service import init_upload, get_filename
+
+    r = asyncio.run(init_upload("my_model.zip", 500))
+    uid = r["upload_id"]
+
+    result = asyncio.run(get_filename(uid))
+
+    assert result == "my_model.zip"
+
+
+def test_sweep_removes_old_uploads(tmp_path, monkeypatch):
+    import asyncio
+    from transformerlab.services import upload_service as svc
+
+    monkeypatch.setattr(svc, "STAGING_ROOT", str(tmp_path / "staging"))
+
+    r = asyncio.run(svc.init_upload("old.bin", 10))
+    uid = r["upload_id"]
+
+    # Backdating meta.json created_at by 25 hours
+    from datetime import datetime, timezone, timedelta
+
+    meta_path = os.path.join(svc.STAGING_ROOT, uid, "meta.json")
+    with open(meta_path) as f:
+        meta = json.load(f)
+    meta["created_at"] = (datetime.now(timezone.utc) - timedelta(hours=25)).isoformat()
+    with open(meta_path, "w") as f:
+        json.dump(meta, f)
+
+    count = svc.sweep_expired_uploads(max_age_hours=24)
+
+    assert count == 1
+    assert not os.path.isdir(os.path.join(svc.STAGING_ROOT, uid))

api/transformerlab/routers/compute_provider/launch.py

@@ -1,5 +1,7 @@
 """Sub-router for compute provider launch logic, including sweep dispatch."""
 
+from typing import Optional
+
 from fastapi import APIRouter, Depends, File, UploadFile
 from sqlalchemy.ext.asyncio import AsyncSession
 from transformerlab.shared.models.user_model import get_async_session
@@ -22,13 +24,14 @@
 async def upload_task_file_for_provider(
     provider_id: str,
     task_id: str,
-    file: UploadFile = File(...),
+    upload_id: Optional[str] = None,
+    file: Optional[UploadFile] = File(None),
     user_and_team=Depends(get_user_and_team),
     session: AsyncSession = Depends(get_async_session),
 ):
     """Upload a single file for a provider-backed task."""
     team_id = user_and_team["team_id"]
-    return await upload_task_file_service(session, team_id, provider_id, task_id, file)
+    return await upload_task_file_service(session, team_id, provider_id, task_id, file, upload_id)
 
 
 @router.post("/")