Merge pull request #660 from transformerlab/add/use-fsspec

Use fsspec in everything

Author: deep1401

api.py

@@ -78,6 +78,7 @@
 )
 from transformerlab.shared.request_context import set_current_org_id
 from lab.dirs import set_organization_id as lab_set_org_id
+from lab import storage
 
 from dotenv import load_dotenv
 
@@ -91,7 +92,7 @@
 # to be overriden by the user.
 os.environ["_TFL_SOURCE_CODE_DIR"] = dirs.TFL_SOURCE_CODE_DIR
 # The temporary image directory for transformerlab (default; per-request overrides computed in routes)
-temp_image_dir = os.path.join(get_workspace_dir(), "temp", "images")
+temp_image_dir = storage.join(get_workspace_dir(), "temp", "images")
 os.environ["TLAB_TEMP_IMAGE_DIR"] = str(temp_image_dir)
 
 # Check if anything is stored in GPU_ORCHESTRATION_SERVER, and if so, print a message
@@ -243,7 +244,10 @@ async def validation_exception_handler(request, exc):
 
 def spawn_fastchat_controller_subprocess():
     global controller_process
-    logfile = open(os.path.join(dirs.FASTCHAT_LOGS_DIR, "controller.log"), "w")
+    controller_log_path = storage.join(dirs.FASTCHAT_LOGS_DIR, "controller.log")
+    # Note: subprocess requires a local file handle, so we use open() directly
+    # but construct the path using storage.join for workspace consistency
+    logfile = open(controller_log_path, "w")
     port = "21001"
 
     controller_process = subprocess.Popen(
@@ -254,7 +258,7 @@ def spawn_fastchat_controller_subprocess():
             "--port",
             port,
             "--log-file",
-            os.path.join(dirs.FASTCHAT_LOGS_DIR, "controller.log"),
+            controller_log_path,
         ],
         stdout=logfile,
         stderr=logfile,
@@ -379,7 +383,7 @@ async def server_worker_start(
 
     from lab.dirs import get_global_log_path
 
-    with open(get_global_log_path(), "a") as global_log:
+    with storage.open(get_global_log_path(), "a") as global_log:
         global_log.write(f"🏃 Loading Inference Server for {model_name} with {inference_params}\n")
 
     process = await shared.async_run_python_daemon_and_update_status(
@@ -392,7 +396,7 @@ async def server_worker_start(
     if exitcode == 99:
         from lab.dirs import get_global_log_path
 
-        with open(get_global_log_path(), "a") as global_log:
+        with storage.open(get_global_log_path(), "a") as global_log:
             global_log.write(
                 "GPU (CUDA) Out of Memory: Please try a smaller model or a different inference engine. Restarting the server may free up resources.\n"
             )
@@ -403,7 +407,7 @@ async def server_worker_start(
     if exitcode is not None and exitcode != 0:
         from lab.dirs import get_global_log_path
 
-        with open(get_global_log_path(), "a") as global_log:
+        with storage.open(get_global_log_path(), "a") as global_log:
             global_log.write(f"Error loading model: {model_name} with exit code {exitcode}\n")
         job = job_get(job_id)
         error_msg = None
@@ -415,7 +419,7 @@ async def server_worker_start(
         return {"status": "error", "message": error_msg}
     from lab.dirs import get_global_log_path
 
-    with open(get_global_log_path(), "a") as global_log:
+    with storage.open(get_global_log_path(), "a") as global_log:
         global_log.write(f"Model loaded successfully: {model_name}\n")
     return {"status": "success", "job_id": job_id}
 

requirements-no-gpu-uv.txt

@@ -556,7 +556,7 @@ tqdm==4.66.5
     #   peft
     #   sentence-transformers
     #   transformers
-transformerlab==0.0.44
+transformerlab==0.0.45
     # via
     #   -r requirements.in
     #   transformerlab-inference

requirements-rocm-uv.txt

@@ -559,7 +559,7 @@ tqdm==4.66.5
     #   peft
     #   sentence-transformers
     #   transformers
-transformerlab==0.0.44
+transformerlab==0.0.45
     # via
     #   -r requirements-rocm.in
     #   transformerlab-inference

requirements-rocm.in

@@ -32,7 +32,7 @@ hf_xet
 macmon-python
 mcp[cli]
 transformerlab-inference==0.2.51
-transformerlab==0.0.44
+transformerlab==0.0.45
 diffusers==0.33.1
 pyrsmi
 controlnet_aux==0.0.10

requirements-uv.txt

@@ -594,7 +594,7 @@ tqdm==4.66.5
     #   peft
     #   sentence-transformers
     #   transformers
-transformerlab==0.0.44
+transformerlab==0.0.45
     # via
     #   -r requirements.in
     #   transformerlab-inference

requirements.in

@@ -31,7 +31,7 @@ markitdown[all]
 hf_xet
 macmon-python
 transformerlab-inference==0.2.51
-transformerlab==0.0.44
+transformerlab==0.0.45
 diffusers==0.33.1
 nvidia-ml-py
 mcp[cli]

test/api/test_diffusion.py

@@ -252,8 +252,8 @@ def test_get_image_by_id_index_out_of_range(client):
     with (
         patch("transformerlab.routers.experiment.diffusion.find_image_by_id") as mock_find_image,
         patch("transformerlab.routers.experiment.diffusion.get_images_dir", return_value="/fake/images"),
-        patch("os.path.exists", return_value=True),
-        patch("os.path.isdir", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.isdir", return_value=True),
     ):
         # Create mock image with folder-format
         mock_image = MagicMock()
@@ -272,9 +272,9 @@ def test_get_image_info_by_id_success(client):
     """Test getting image metadata by ID"""
     with (
         patch("transformerlab.routers.experiment.diffusion.find_image_by_id") as mock_find_image,
-        patch("os.path.exists", return_value=True),
-        patch("os.path.isdir", return_value=True),
-        patch("os.listdir", return_value=["0.png", "1.png", "2.png"]),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.isdir", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.ls", return_value=["0.png", "1.png", "2.png"]),
     ):
         # Create mock image
         mock_image = MagicMock()
@@ -295,9 +295,9 @@ def test_get_image_count_success(client):
     """Test getting image count for an image set"""
     with (
         patch("transformerlab.routers.experiment.diffusion.find_image_by_id") as mock_find_image,
-        patch("os.path.exists", return_value=True),
-        patch("os.path.isdir", return_value=True),
-        patch("os.listdir", return_value=["0.png", "1.png"]),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.isdir", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.ls", return_value=["0.png", "1.png"]),
     ):
         # Create mock image
         mock_image = MagicMock()
@@ -317,8 +317,8 @@ def test_delete_image_from_history_not_found(client):
     """Test deleting a non-existent image from history"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data='[{"id": "other-id", "image_path": "/fake/path.png"}]')),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data='[{"id": "other-id", "image_path": "/fake/path.png"}]')),
     ):
         resp = client.delete("/experiment/test-exp-name/diffusion/history/non-existent-id")
         assert resp.status_code == 500
@@ -331,12 +331,12 @@ def test_create_dataset_from_history_success(client):
         patch("transformerlab.routers.experiment.diffusion.find_image_by_id") as mock_find_image,
         patch("transformerlab.routers.experiment.diffusion.Dataset.get") as mock_dataset_get,
         patch("transformerlab.routers.experiment.diffusion.create_local_dataset") as mock_create_dataset,
-        patch("os.makedirs"),
-        patch("os.path.exists", return_value=True),
-        patch("os.path.isdir", return_value=True),
-        patch("os.listdir", return_value=["0.png", "1.png"]),
-        patch("shutil.copy2"),
-        patch("builtins.open", mock_open()),
+        patch("transformerlab.routers.experiment.diffusion.storage.makedirs"),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.isdir", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.ls", return_value=["/fake/path/folder/0.png", "/fake/path/folder/1.png"]),
+        patch("transformerlab.routers.experiment.diffusion.storage.copy_file"),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open()),
     ):
         # Mock Dataset.get to raise FileNotFoundError for non-existent dataset (new behavior)
         mock_dataset_get.side_effect = FileNotFoundError("Directory for Dataset with id 'test-dataset' not found")
@@ -610,9 +610,9 @@ def test_load_history_success():
     """Test loading history with valid data"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
         patch(
-            "builtins.open",
+            "transformerlab.routers.experiment.diffusion.storage.open",
             mock_open(
                 read_data='[{"id": "test-id", "model": "test-model", "prompt": "test prompt", "adaptor": "", "adaptor_scale": 1.0, "num_inference_steps": 20, "guidance_scale": 7.5, "seed": 42, "image_path": "/fake/path.png", "timestamp": "2023-01-01T00:00:00", "upscaled": false, "upscale_factor": 1, "negative_prompt": "", "eta": 0.0, "clip_skip": 0, "guidance_rescale": 0.0, "height": 512, "width": 512, "generation_time": 5.0, "num_images": 1, "input_image_path": "", "strength": 0.8, "is_img2img": false, "mask_image_path": "", "is_inpainting": false}]'
             ),
@@ -665,8 +665,8 @@ def test_load_history_with_pagination():
 
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data=json.dumps(history_data))),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data=json.dumps(history_data))),
     ):
         from transformerlab.routers.experiment.diffusion import load_history
 
@@ -683,7 +683,7 @@ def test_load_history_no_file():
     """Test loading history when history file doesn't exist"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=False),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=False),
     ):
         from transformerlab.routers.experiment.diffusion import load_history
 
@@ -697,8 +697,8 @@ def test_load_history_invalid_json():
     """Test loading history with corrupted JSON file"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data="invalid json")),
+        patch("lab.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data="invalid json")),
     ):
         from transformerlab.routers.experiment.diffusion import load_history
 
@@ -769,8 +769,8 @@ def test_find_image_by_id_success():
 
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data=json.dumps(history_data))),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data=json.dumps(history_data))),
     ):
         from transformerlab.routers.experiment.diffusion import find_image_by_id
 
@@ -816,8 +816,8 @@ def test_find_image_by_id_not_found(client):
 
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data=json.dumps(history_data))),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data=json.dumps(history_data))),
     ):
         from transformerlab.routers.experiment.diffusion import find_image_by_id
 
@@ -830,7 +830,7 @@ def test_find_image_by_id_no_file():
     """Test finding an image by ID when history file doesn't exist"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=False),
+        patch("transformerlab.routers.experiment.diffusion.storage.exists", return_value=False),
     ):
         from transformerlab.routers.experiment.diffusion import find_image_by_id
 
@@ -843,8 +843,8 @@ def test_find_image_by_id_invalid_json():
     """Test finding an image by ID with corrupted JSON file"""
     with (
         patch("transformerlab.routers.experiment.diffusion.get_history_file_path", return_value="/fake/history.json"),
-        patch("os.path.exists", return_value=True),
-        patch("builtins.open", mock_open(read_data="invalid json")),
+        patch("lab.storage.exists", return_value=True),
+        patch("transformerlab.routers.experiment.diffusion.storage.open", mock_open(read_data="invalid json")),
     ):
         from transformerlab.routers.experiment.diffusion import find_image_by_id
 

test/api/test_experiment_export.py

@@ -3,6 +3,7 @@
 from transformerlab.services import experiment_service
 from transformerlab.services.tasks_service import tasks_service
 from lab.dirs import get_workspace_dir
+from lab import storage
 
 
 async def test_export_experiment(client):
@@ -66,13 +67,13 @@ async def test_export_experiment(client):
     # The response should be a JSON file
     assert response.headers["content-type"] == "application/json"
 
-    WORKSPACE_DIR = get_workspace_dir()
+    workspace_dir = get_workspace_dir()
 
     # Read the exported file from workspace directory
-    export_file = os.path.join(WORKSPACE_DIR, f"{test_experiment_name}_export.json")
-    assert os.path.exists(export_file)
+    export_file = storage.join(workspace_dir, f"{test_experiment_name}_export.json")
+    assert storage.exists(export_file)
 
-    with open(export_file, "r") as f:
+    with storage.open(export_file, "r") as f:
         exported_data = json.load(f)
 
     # Verify the exported data structure
@@ -104,5 +105,5 @@ async def test_export_experiment(client):
 
     # Clean up
     experiment_service.experiment_delete(experiment_id)
-    if os.path.exists(export_file):
-        os.remove(export_file)
+    if storage.exists(export_file):
+        storage.rm(export_file)

transformerlab/db/constants.py

@@ -1,7 +1,6 @@
 # --- Centralized Database Configuration ---
 from lab import HOME_DIR
 
-
 db = None  # This will hold the aiosqlite connection
 DATABASE_FILE_NAME = f"{HOME_DIR}/llmlab.sqlite3"
 DATABASE_URL = f"sqlite+aiosqlite:///{DATABASE_FILE_NAME}"

transformerlab/fastchat_openai_api.py

@@ -45,7 +45,7 @@
     UsageInfo,
 )
 from pydantic import BaseModel as PydanticBaseModel
-from lab import Experiment
+from lab import Experiment, storage
 
 WORKER_API_TIMEOUT = 3600
 
@@ -233,16 +233,23 @@ def log_prompt(prompt):
     from lab.dirs import get_logs_dir
 
     logs_dir = get_logs_dir()
-    if os.path.exists(os.path.join(logs_dir, "prompt.log")):
-        if os.path.getsize(os.path.join(logs_dir, "prompt.log")) > MAX_LOG_SIZE_BEFORE_ROTATE:
-            with open(os.path.join(logs_dir, "prompt.log"), "r") as f:
+    prompt_log_path = storage.join(logs_dir, "prompt.log")
+    if storage.exists(prompt_log_path):
+        # Get file size - for remote storage, we may need to read the file to check size
+        try:
+            with storage.open(prompt_log_path, "r") as f:
                 lines = f.readlines()
-            with open(os.path.join(logs_dir, "prompt.log"), "w") as f:
-                f.writelines(lines[-1000:])
-            with open(os.path.join(logs_dir, f"prompt_{time.strftime('%Y%m%d%H%M%S')}.log"), "w") as f:
-                f.writelines(lines[:-1000])
-
-    with open(os.path.join(logs_dir, "prompt.log"), "a") as f:
+            file_size = sum(len(line.encode('utf-8')) for line in lines)
+            if file_size > MAX_LOG_SIZE_BEFORE_ROTATE:
+                with storage.open(prompt_log_path, "w") as f:
+                    f.writelines(lines[-1000:])
+                with storage.open(storage.join(logs_dir, f"prompt_{time.strftime('%Y%m%d%H%M%S')}.log"), "w") as f:
+                    f.writelines(lines[:-1000])
+        except Exception:
+            # If we can't read the file, just continue with appending
+            pass
+
+    with storage.open(prompt_log_path, "a") as f:
         log_entry = {}
         log_entry["date"] = time.strftime("%Y-%m-%d %H:%M:%S")
         log_entry["log"] = prompt
@@ -254,7 +261,10 @@ def log_prompt(prompt):
 async def get_prompt_log():
     from lab.dirs import get_logs_dir
 
-    return FileResponse(os.path.join(get_logs_dir(), "prompt.log"))
+    prompt_log_path = storage.join(get_logs_dir(), "prompt.log")
+    # FileResponse needs a local file path, so use the path string directly
+    # For remote storage, this would need special handling
+    return FileResponse(prompt_log_path)
 
 
 async def check_length(request, prompt, max_tokens):
@@ -526,8 +536,8 @@ async def create_audio_tts(request: AudioGenerationRequest):
     exp_obj = Experiment.get(request.experiment_id)
     experiment_dir = exp_obj.get_dir()
 
-    audio_dir = os.path.join(experiment_dir, "audio")
-    os.makedirs(audio_dir, exist_ok=True)
+    audio_dir = storage.join(experiment_dir, "audio")
+    storage.makedirs(audio_dir, exist_ok=True)
 
     gen_params = {
         "audio_dir": audio_dir,
@@ -560,16 +570,16 @@ async def create_audio_tts(request: AudioGenerationRequest):
 async def upload_audio_reference(experimentId: str, audio: UploadFile = File(...)):
     exp_obj = Experiment(experimentId)
     experiment_dir = exp_obj.get_dir()
-    uploaded_audio_dir = os.path.join(experiment_dir, "uploaded_audio")
-    os.makedirs(uploaded_audio_dir, exist_ok=True)
+    uploaded_audio_dir = storage.join(experiment_dir, "uploaded_audio")
+    storage.makedirs(uploaded_audio_dir, exist_ok=True)
 
     file_prefix = str(uuid.uuid4())
     _, ext = os.path.splitext(audio.filename)
-    file_path = os.path.join(uploaded_audio_dir, file_prefix + ext)
+    file_path = storage.join(uploaded_audio_dir, file_prefix + ext)
 
     # Save the uploaded file
-    with open(file_path, "wb") as f:
-        content = await audio.read()
+    content = await audio.read()
+    with storage.open(file_path, "wb") as f:
         f.write(content)
 
     return JSONResponse({"audioPath": file_path})