Merge pull request #1111 from transformerlab/add/fsspec-async

Fsspec async

Author: deep1401

.github/workflows/pytest-sdk.yml

@@ -41,7 +41,5 @@ jobs:
           uv pip install -e .
 
       - name: Run SDK tests (uv)
-        env:
-          PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1"
         run: |
           uv run pytest

.github/workflows/pytest.yml

@@ -51,7 +51,6 @@ jobs:
       - name: Test with pytest
         run: |
           pytest --cov=transformerlab --cov-branch --cov-report=xml -k 'not test_teams'
-
       - name: Upload results to Codecov
         uses: codecov/codecov-action@v5
         with:

api/api.py

@@ -104,28 +104,33 @@
 # used internally to set constants that are shared between separate processes. They are not meant to be
 # 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 = storage.join(get_workspace_dir(), "temp", "images")
-os.environ["TLAB_TEMP_IMAGE_DIR"] = str(temp_image_dir)
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Docs on lifespan events: https://fastapi.tiangolo.com/advanced/events/"""
     # Do the following at API Startup:
     print_launch_message()
+    # Initialize directories early
+    from transformerlab.shared import dirs as shared_dirs
+
+    await shared_dirs.initialize_dirs()
+
+    # Set the temporary image directory for transformerlab (computed async)
+    temp_image_dir = storage.join(await get_workspace_dir(), "temp", "images")
+    os.environ["TLAB_TEMP_IMAGE_DIR"] = str(temp_image_dir)
     # Validate cloud credentials early - fail fast if missing
     validate_cloud_credentials()
-    galleries.update_gallery_cache()
+    await galleries.update_gallery_cache()
     spawn_fastchat_controller_subprocess()
     await db.init()  # This now runs Alembic migrations internally
     print("✅ SEED DATA")
     # Initialize experiments
-    seed_default_experiments()
+    await seed_default_experiments()
     # Seed default admin user
     await seed_default_admin_user()
     # Cancel any running jobs
-    cancel_in_progress_jobs()
+    await cancel_in_progress_jobs()
 
     # Create buckets for all existing teams if TFL_API_STORAGE_URI is enabled
     if os.getenv("TFL_API_STORAGE_URI"):
@@ -362,7 +367,7 @@ async def server_worker_start(
     # then we check to see if we are an experiment
     elif experiment_id is not None:
         try:
-            experiment = experiment_get(experiment_id)
+            experiment = await experiment_get(experiment_id)
             experiment_config = (
                 experiment["config"]
                 if isinstance(experiment["config"], dict)
@@ -394,15 +399,15 @@ async def server_worker_start(
     model_architecture = model_architecture
 
     plugin_name = inference_engine
-    plugin_location = lab_dirs.plugin_dir_by_name(plugin_name)
+    plugin_location = await lab_dirs.plugin_dir_by_name(plugin_name)
 
     model = model_name
     if model_filename is not None and model_filename != "":
         model = model_filename
 
     if adaptor != "":
         # Resolve per-request workspace if multitenant
-        workspace_dir = get_workspace_dir()
+        workspace_dir = await get_workspace_dir()
         adaptor = f"{workspace_dir}/adaptors/{secure_filename(model)}/{adaptor}"
 
     params = [
@@ -419,14 +424,14 @@ async def server_worker_start(
         json.dumps(inference_params),
     ]
 
-    job_id = job_create(type="LOAD_MODEL", status="STARTED", job_data="{}", experiment_id=experiment_id)
+    job_id = await job_create(type="LOAD_MODEL", status="STARTED", job_data="{}", experiment_id=experiment_id)
 
     print("Loading plugin loader instead of default worker")
 
     from lab.dirs import get_global_log_path
 
-    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")
+    async with await storage.open(await get_global_log_path(), "a") as global_log:
+        await global_log.write(f"🏃 Loading Inference Server for {model_name} with {inference_params}\n")
 
     # Pass organization_id as environment variable to subprocess
     from transformerlab.shared.request_context import get_current_org_id
@@ -447,8 +452,8 @@ async def server_worker_start(
     if exitcode == 99:
         from lab.dirs import get_global_log_path
 
-        with storage.open(get_global_log_path(), "a") as global_log:
-            global_log.write(
+        async with await storage.open(await get_global_log_path(), "a") as global_log:
+            await 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"
             )
         return {
@@ -458,20 +463,20 @@ async def server_worker_start(
     if exitcode is not None and exitcode != 0:
         from lab.dirs import get_global_log_path
 
-        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)
+        async with await storage.open(await get_global_log_path(), "a") as global_log:
+            await global_log.write(f"Error loading model: {model_name} with exit code {exitcode}\n")
+        job = await job_get(job_id)
         error_msg = None
         if job and job.get("job_data"):
             error_msg = job["job_data"].get("error_msg")
         if not error_msg:
             error_msg = f"Exit code {exitcode}"
-            job_update_status(job_id, "FAILED", experiment_id=experiment_id, error_msg=error_msg)
+            await job_update_status(job_id, "FAILED", experiment_id=experiment_id, error_msg=error_msg)
         return {"status": "error", "message": error_msg}
     from lab.dirs import get_global_log_path
 
-    with storage.open(get_global_log_path(), "a") as global_log:
-        global_log.write(f"Model loaded successfully: {model_name}\n")
+    async with await storage.open(await get_global_log_path(), "a") as global_log:
+        await global_log.write(f"Model loaded successfully: {model_name}\n")
     return {"status": "success", "job_id": job_id}
 
 
@@ -630,7 +635,9 @@ def run():
     )
 
     if args.https:
-        cert_path, key_path = ensure_persistent_self_signed_cert()
+        import asyncio
+
+        cert_path, key_path = asyncio.run(ensure_persistent_self_signed_cert())
         uvicorn.run(
             "api:app", host=args.host, port=args.port, log_level="warning", ssl_certfile=cert_path, ssl_keyfile=key_path
         )

api/pyproject.toml

@@ -33,8 +33,8 @@ dependencies = [
   "soundfile==0.13.1",
   "tensorboardX==2.6.2.2",
   "timm==1.0.15",
-  "transformerlab==0.0.58",
-  "transformerlab-inference==0.2.51",
+  "transformerlab==0.0.62",
+  "transformerlab-inference==0.2.52",
   "transformers==4.57.1",
   "wandb==0.19.10",
   "werkzeug==3.1.3",
@@ -109,4 +109,4 @@ cpu = [
   "tensorboard==2.18.0",
   "tiktoken==0.8.0",
   "watchfiles==1.0.4",
-]
+]

api/test/api/conftest.py

@@ -5,6 +5,7 @@
 
 # Create test directories before setting environment variables
 os.makedirs("test/tmp/", exist_ok=True)
+os.makedirs("test/tmp/webapp", exist_ok=True)  # Create webapp directory for static files
 
 os.environ["TFL_HOME_DIR"] = "test/tmp/"
 # Note: TFL_WORKSPACE_DIR is not set so that get_workspace_dir() will use the org-based

api/test/api/test_data.py

@@ -1,6 +1,7 @@
 import pytest
 import os
 import json
+import asyncio
 from io import BytesIO
 from PIL import Image
 from pathlib import Path
@@ -13,7 +14,7 @@ def cleanup_dataset(dataset_id, client):
     from transformerlab.shared.shared import slugify
     import shutil
 
-    dataset_dir = dirs.dataset_dir_by_id(slugify(dataset_id))
+    dataset_dir = asyncio.run(dirs.dataset_dir_by_id(slugify(dataset_id)))
     shutil.rmtree(dataset_dir, ignore_errors=True)
     client.get(f"/data/delete?dataset_id={dataset_id}")
 
@@ -54,7 +55,7 @@ def test_data_info(client):
 def test_save_metadata(client):
     source_dataset_id = "source_dataset"
     new_dataset_id = "destination_dataset"
-    dataset_dir = dirs.dataset_dir_by_id(slugify(source_dataset_id))
+    dataset_dir = asyncio.run(dirs.dataset_dir_by_id(slugify(source_dataset_id)))
     os.makedirs(dataset_dir, exist_ok=True)
 
     # Create dummy JPEG image
@@ -96,7 +97,7 @@ def test_save_metadata(client):
     data = response.json()
     assert data["status"] == "success"
 
-    new_dataset_dir = Path(dirs.dataset_dir_by_id(slugify(new_dataset_id)))
+    new_dataset_dir = Path(asyncio.run(dirs.dataset_dir_by_id(slugify(new_dataset_id))))
     assert new_dataset_dir.exists()
 
     cleanup_dataset(source_dataset_id, client)
@@ -106,7 +107,7 @@ def test_save_metadata(client):
 @pytest.mark.skip(reason="Skipping as it contains application-specific logic")
 def test_edit_with_template(client):
     dataset_id = "test_dataset"
-    dataset_dir = dirs.dataset_dir_by_id(slugify(dataset_id))
+    dataset_dir = asyncio.run(dirs.dataset_dir_by_id(slugify(dataset_id)))
     os.makedirs(dataset_dir, exist_ok=True)
 
     image_path = os.path.join(dataset_dir, "image.jpg")

api/test/api/test_dataset_service.py

@@ -16,7 +16,8 @@ def tmp_dataset_dir(tmp_path: Path) -> Path:
     return tmp_path
 
 
-def test_load_local_dataset_filters_index_and_hidden(tmp_dataset_dir: Path, monkeypatch):
+@pytest.mark.asyncio
+async def test_load_local_dataset_filters_index_and_hidden(tmp_dataset_dir: Path, monkeypatch):
     # Import inside test to ensure module path resolution for monkeypatching
     from transformerlab.services import dataset_service
 
@@ -30,7 +31,7 @@ def fake_load_dataset(path=None, data_files=None, streaming=False):
 
     monkeypatch.setattr(dataset_service, "load_dataset", fake_load_dataset)
 
-    result = dataset_service.load_local_dataset(str(tmp_dataset_dir))
+    result = await dataset_service.load_local_dataset(str(tmp_dataset_dir))
 
     assert result == {"ok": True}
     assert captured["path"] == str(tmp_dataset_dir)
@@ -43,7 +44,8 @@ def fake_load_dataset(path=None, data_files=None, streaming=False):
     assert captured["streaming"] is False
 
 
-def test_load_local_dataset_uses_explicit_data_files(tmp_path: Path, monkeypatch):
+@pytest.mark.asyncio
+async def test_load_local_dataset_uses_explicit_data_files(tmp_path: Path, monkeypatch):
     from transformerlab.services import dataset_service
 
     # Explicit files provided (note: function should not re-filter these)
@@ -60,7 +62,9 @@ def fake_load_dataset(path=None, data_files=None, streaming=False):
 
     monkeypatch.setattr(dataset_service, "load_dataset", fake_load_dataset)
 
-    result = dataset_service.load_local_dataset(str(tmp_path), data_files=["keep.me", "index.json"], streaming=True)
+    result = await dataset_service.load_local_dataset(
+        str(tmp_path), data_files=["keep.me", "index.json"], streaming=True
+    )
 
     assert result == {"ok": True}
     assert captured["path"] == str(tmp_path)
@@ -72,7 +76,8 @@ def fake_load_dataset(path=None, data_files=None, streaming=False):
     assert captured["streaming"] is True
 
 
-def test_load_local_dataset_fallback_when_no_valid_files(tmp_path: Path, monkeypatch):
+@pytest.mark.asyncio
+async def test_load_local_dataset_fallback_when_no_valid_files(tmp_path: Path, monkeypatch):
     from transformerlab.services import dataset_service
 
     # Only metadata/hidden files present
@@ -89,7 +94,7 @@ def fake_load_dataset(path=None, data_files=None, streaming=False):
 
     monkeypatch.setattr(dataset_service, "load_dataset", fake_load_dataset)
 
-    result = dataset_service.load_local_dataset(str(tmp_path))
+    result = await dataset_service.load_local_dataset(str(tmp_path))
 
     assert result == {"ok": True}
     assert captured["path"] == str(tmp_path)