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

Switch to using fsspec everywhere

Author: deep1401

.gitignore

@@ -205,3 +205,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+
+scripts/examples/test_fsspec.sh

pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab"
-version = "0.0.44"
+version = "0.0.45"
 description = "Python SDK for Transformer Lab"
 readme = "README.md"
 requires-python = ">=3.10"
 authors = [{ name = "Transformer Lab", email = "developers@transformerlab.ai" }]
 license = { file = "LICENSE" }
-dependencies = ["werkzeug", "pytest"]
+dependencies = ["werkzeug", "pytest", "wandb", "fsspec", "s3fs"]
 
 [project.urls]
 "Homepage" = "https://github.com/transformerlab/transformerlab-sdk"

src/lab/dataset.py

@@ -1,15 +1,15 @@
-import os
 from werkzeug.utils import secure_filename
 
 from .dirs import get_datasets_dir
 from .labresource import BaseLabResource
+from . import storage
 
 
 class Dataset(BaseLabResource):
     def get_dir(self):
         """Abstract method on BaseLabResource"""
         dataset_id_safe = secure_filename(str(self.id))
-        return os.path.join(get_datasets_dir(), dataset_id_safe)
+        return storage.join(get_datasets_dir(), dataset_id_safe)
 
     def _default_json(self):
         # Default metadata modeled after API dataset table fields
@@ -45,14 +45,18 @@ def get_metadata(self):
     def list_all():
         results = []
         datasets_dir = get_datasets_dir()
-        if not os.path.isdir(datasets_dir):
+        if not storage.isdir(datasets_dir):
             return results
-        for entry in os.listdir(datasets_dir):
-            full = os.path.join(datasets_dir, entry)
-            if not os.path.isdir(full):
+        try:
+            entries = storage.ls(datasets_dir, detail=False)
+        except Exception:
+            entries = []
+        for full in entries:
+            if not storage.isdir(full):
                 continue
             # Attempt to read index.json (or latest snapshot)
             try:
+                entry = full.rstrip("/").split("/")[-1]
                 ds = Dataset(entry)
                 results.append(ds.get_metadata())
             except Exception:

src/lab/dirs.py

@@ -3,52 +3,68 @@
 import os
 import contextvars
 from werkzeug.utils import secure_filename
+from . import storage
+from .storage import _current_tfl_storage_uri
 
 # TFL_HOME_DIR
-if "TFL_HOME_DIR" in os.environ:
+if "TFL_HOME_DIR" in os.environ and not (_current_tfl_storage_uri.get() or os.getenv("TFL_STORAGE_URI")):
     HOME_DIR = os.environ["TFL_HOME_DIR"]
     if not os.path.exists(HOME_DIR):
         print(f"Error: Home directory {HOME_DIR} does not exist")
         exit(1)
     print(f"Home directory is set to: {HOME_DIR}")
 else:
-    HOME_DIR = os.path.join(os.path.expanduser("~"), ".transformerlab")
-    os.makedirs(name=HOME_DIR, exist_ok=True)
-    print(f"Using default home directory: {HOME_DIR}")
+    # If TFL_STORAGE_URI is set (via context or env), HOME_DIR concept maps to storage.root_uri()
+    HOME_DIR = storage.root_uri() if (_current_tfl_storage_uri.get() or os.getenv("TFL_STORAGE_URI")) else os.path.join(os.path.expanduser("~"), ".transformerlab")
+    if not (_current_tfl_storage_uri.get() or os.getenv("TFL_STORAGE_URI")):
+        os.makedirs(name=HOME_DIR, exist_ok=True)
+        print(f"Using default home directory: {HOME_DIR}")
 
 # Context var for organization id (set by host app/session)
 _current_org_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
     "current_org_id", default=None
 )
 
-
 def set_organization_id(organization_id: str | None) -> None:
     _current_org_id.set(organization_id)
+    if organization_id is not None:
+        _current_tfl_storage_uri.set(os.getenv("TFL_API_STORAGE_URI"))
+    else:
+        _current_tfl_storage_uri.set(None)
 
 
 def get_workspace_dir() -> str:
     # Remote SkyPilot workspace override (highest precedence)
     # Only return container workspace path when value is exactly "true"
     if os.getenv("_TFL_REMOTE_SKYPILOT_WORKSPACE") == "true":
+        if os.getenv("TFL_STORAGE_URI") is not None:
+            return storage.root_uri()
+        
         return "/workspace"
 
     # Explicit override wins
-    if "TFL_WORKSPACE_DIR" in os.environ:
+    if "TFL_WORKSPACE_DIR" in os.environ and not (_current_tfl_storage_uri.get() is not None and os.getenv("TFL_STORAGE_URI") is not None):
         value = os.environ["TFL_WORKSPACE_DIR"]
         if not os.path.exists(value):
             print(f"Error: Workspace directory {value} does not exist")
             exit(1)
         return value
 
     org_id = _current_org_id.get()
+
     if org_id:
-        path = os.path.join(HOME_DIR, "orgs", org_id, "workspace")
-        os.makedirs(name=path, exist_ok=True)
+        # If the storage URI is set, use it for the org workspace
+        if _current_tfl_storage_uri.get() is not None:
+            return _current_tfl_storage_uri.get()
+        path = storage.join(HOME_DIR, "orgs", org_id, "workspace")
+        storage.makedirs(path, exist_ok=True)
         return path
+    
+    if os.getenv("TFL_STORAGE_URI"):
+        return storage.root_uri()
 
-    # Default single-tenant path
-    path = os.path.join(HOME_DIR, "workspace")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(HOME_DIR, "workspace")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
@@ -79,92 +95,96 @@ def get_workspace_dir() -> str:
 
 
 def get_experiments_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "experiments")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "experiments")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_jobs_dir() -> str:
     workspace_dir = get_workspace_dir()
-    path = os.path.join(workspace_dir, "jobs")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(workspace_dir, "jobs")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_global_log_path() -> str:
-    return os.path.join(get_workspace_dir(), "transformerlab.log")
+    return storage.join(get_workspace_dir(), "transformerlab.log")
 
 
 def get_logs_dir() -> str:
-    path = os.path.join(HOME_DIR, "logs")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(HOME_DIR, "logs")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 # TODO: Move this to Experiment
 def experiment_dir_by_name(experiment_name: str) -> str:
     experiments_dir = get_experiments_dir()
-    return os.path.join(experiments_dir, experiment_name)
+    return storage.join(experiments_dir, experiment_name)
 
 
 def get_plugin_dir() -> str:
-    return os.path.join(get_workspace_dir(), "plugins")
+    return storage.join(get_workspace_dir(), "plugins")
 
 
 def plugin_dir_by_name(plugin_name: str) -> str:
     plugin_name = secure_filename(plugin_name)
-    return os.path.join(get_plugin_dir(), plugin_name)
+    return storage.join(get_plugin_dir(), plugin_name)
 
 
 def get_models_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "models")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "models")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_datasets_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "datasets")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "datasets")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_tasks_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "tasks")
-    os.makedirs(name=path, exist_ok=True)
+    tfl_storage_uri = _current_tfl_storage_uri.get()
+    if tfl_storage_uri is not None:
+        return storage.join(tfl_storage_uri, "tasks")
+
+    path = storage.join(get_workspace_dir(), "tasks")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def dataset_dir_by_id(dataset_id: str) -> str:
-    return os.path.join(get_datasets_dir(), dataset_id)
+    return storage.join(get_datasets_dir(), dataset_id)
 
 
 def get_temp_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "temp")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "temp")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_prompt_templates_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "prompt_templates")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "prompt_templates")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_tools_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "tools")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "tools")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_batched_prompts_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "batched_prompts")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "batched_prompts")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
 def get_galleries_cache_dir() -> str:
-    path = os.path.join(get_workspace_dir(), "galleries")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_workspace_dir(), "galleries")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
@@ -175,16 +195,16 @@ def get_job_dir(job_id: str | int) -> str:
     instance is not readily available.
     """
     job_id_safe = secure_filename(str(job_id))
-    return os.path.join(get_jobs_dir(), job_id_safe)
+    return storage.join(get_jobs_dir(), job_id_safe)
 
 
 def get_job_artifacts_dir(job_id: str | int) -> str:
     """
     Return the artifacts directory for a specific job, creating it if needed.
     Example: ~/.transformerlab/workspace/jobs/<job_id>/artifacts
     """
-    path = os.path.join(get_job_dir(job_id), "artifacts")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_job_dir(job_id), "artifacts")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
@@ -193,8 +213,8 @@ def get_job_checkpoints_dir(job_id: str | int) -> str:
     Return the checkpoints directory for a specific job, creating it if needed.
     Example: ~/.transformerlab/workspace/jobs/<job_id>/checkpoints
     """
-    path = os.path.join(get_job_dir(job_id), "checkpoints")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_job_dir(job_id), "checkpoints")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
@@ -203,8 +223,8 @@ def get_job_eval_results_dir(job_id: str | int) -> str:
     Return the eval_results directory for a specific job, creating it if needed.
     Example: ~/.transformerlab/workspace/jobs/<job_id>/eval_results
     """
-    path = os.path.join(get_job_dir(job_id), "eval_results")
-    os.makedirs(name=path, exist_ok=True)
+    path = storage.join(get_job_dir(job_id), "eval_results")
+    storage.makedirs(path, exist_ok=True)
     return path
 
 
@@ -215,14 +235,14 @@ def get_job_eval_results_dir(job_id: str | int) -> str:
 async def eval_output_file(experiment_name: str, eval_name: str) -> str:
     experiment_dir = experiment_dir_by_name(experiment_name)
     eval_name = secure_filename(eval_name)
-    p = os.path.join(experiment_dir, "evals", eval_name)
-    os.makedirs(p, exist_ok=True)
-    return os.path.join(p, "output.txt")
+    p = storage.join(experiment_dir, "evals", eval_name)
+    storage.makedirs(p, exist_ok=True)
+    return storage.join(p, "output.txt")
 
 
 async def generation_output_file(experiment_name: str, generation_name: str) -> str:
     experiment_dir = experiment_dir_by_name(experiment_name)
     generation_name = secure_filename(generation_name)
-    p = os.path.join(experiment_dir, "generations", generation_name)
-    os.makedirs(p, exist_ok=True)
-    return os.path.join(p, "output.txt")
+    p = storage.join(experiment_dir, "generations", generation_name)
+    storage.makedirs(p, exist_ok=True)
+    return storage.join(p, "output.txt")