Merge pull request #1340 from transformerlab/add/list-datasets-facade

Add an option to access generated datasets via lab facade

Author: deep1401

api/pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
   "soundfile==0.13.1",
   "tensorboardX==2.6.2.2",
   "timm==1.0.15",
-  "transformerlab==0.0.78",
+  "transformerlab==0.0.79",
   "transformerlab-inference==0.2.52",
   "transformers==4.57.1",
   "wandb==0.23.1",

lab-sdk/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab"
-version = "0.0.78"
+version = "0.0.79"
 description = "Python SDK for Transformer Lab"
 readme = "README.md"
 requires-python = ">=3.10"

lab-sdk/src/lab/lab_facade.py

@@ -1519,6 +1519,39 @@ async def async_get_artifact_paths(self) -> list[str]:
         self._ensure_initialized()
         return await self._job.get_artifact_paths()  # type: ignore[union-attr]
 
+    def list_datasets(self) -> list[Dict[str, Any]]:
+        """
+        List all local datasets available in the workspace.
+
+        Returns:
+            List of dictionaries containing dataset metadata. Each dictionary includes:
+            - dataset_id: The dataset identifier
+            - location: Where the dataset is stored (e.g. "local")
+            - description: A human-readable description
+            - size: The dataset size (-1 if unknown)
+            - json_data: Additional dataset metadata
+        """
+        return _run_async(Dataset.list_all())
+
+    def get_dataset(self, dataset_id: str, job_id: Optional[str] = None) -> Dataset:
+        """
+        Get a specific local dataset by ID.
+
+        The lookup is automatically scoped to the current team/organization context.
+
+        Args:
+            dataset_id: The identifier of the dataset to retrieve
+            job_id: Optional job ID. If provided, looks up the dataset in the
+                    job-specific datasets directory instead of the global one.
+
+        Returns:
+            Dataset: A Dataset instance for the specified dataset
+
+        Raises:
+            FileNotFoundError: If the dataset directory doesn't exist
+        """
+        return _run_async(Dataset.get(dataset_id, job_id=job_id))
+
     def list_models(self) -> list[Dict[str, Any]]:
         """
         List all local models available in the workspace.

lab-sdk/tests/test_lab_facade.py

@@ -811,6 +811,97 @@ def test_lab_ensure_initialized(tmp_path, monkeypatch):
             pass  # Expected
 
 
+def test_lab_list_datasets(tmp_path, monkeypatch):
+    _fresh(monkeypatch)
+    home = tmp_path / ".tfl_home"
+    ws = tmp_path / ".tfl_ws"
+    home.mkdir()
+    ws.mkdir()
+    monkeypatch.setenv("TFL_HOME_DIR", str(home))
+    monkeypatch.setenv("TFL_WORKSPACE_DIR", str(ws))
+
+    from lab.lab_facade import Lab
+    from lab.dataset import Dataset
+
+    # Create test datasets
+    ds1 = asyncio.run(Dataset.create("test_dataset_1"))
+    asyncio.run(ds1.set_metadata(description="First dataset"))
+
+    ds2 = asyncio.run(Dataset.create("test_dataset_2"))
+    asyncio.run(ds2.set_metadata(description="Second dataset"))
+
+    lab = Lab()
+    # list_datasets doesn't require initialization
+    datasets = lab.list_datasets()
+
+    assert len(datasets) >= 2
+    dataset_ids = [d.get("dataset_id") for d in datasets]
+    assert "test_dataset_1" in dataset_ids
+    assert "test_dataset_2" in dataset_ids
+
+
+def test_lab_list_datasets_empty(tmp_path, monkeypatch):
+    _fresh(monkeypatch)
+    home = tmp_path / ".tfl_home"
+    ws = tmp_path / ".tfl_ws"
+    home.mkdir()
+    ws.mkdir()
+    monkeypatch.setenv("TFL_HOME_DIR", str(home))
+    monkeypatch.setenv("TFL_WORKSPACE_DIR", str(ws))
+
+    from lab.lab_facade import Lab
+
+    lab = Lab()
+    datasets = lab.list_datasets()
+
+    assert isinstance(datasets, list)
+    assert len(datasets) == 0
+
+
+def test_lab_get_dataset(tmp_path, monkeypatch):
+    _fresh(monkeypatch)
+    home = tmp_path / ".tfl_home"
+    ws = tmp_path / ".tfl_ws"
+    home.mkdir()
+    ws.mkdir()
+    monkeypatch.setenv("TFL_HOME_DIR", str(home))
+    monkeypatch.setenv("TFL_WORKSPACE_DIR", str(ws))
+
+    from lab.lab_facade import Lab
+    from lab.dataset import Dataset
+
+    # Create a test dataset
+    ds = asyncio.run(Dataset.create("test_dataset_get"))
+    asyncio.run(ds.set_metadata(description="My Dataset"))
+
+    lab = Lab()
+    # get_dataset doesn't require initialization
+    retrieved_ds = lab.get_dataset("test_dataset_get")
+
+    assert retrieved_ds.id == "test_dataset_get"
+    metadata = asyncio.run(retrieved_ds.get_metadata())
+    assert metadata["description"] == "My Dataset"
+
+
+def test_lab_get_dataset_nonexistent(tmp_path, monkeypatch):
+    _fresh(monkeypatch)
+    home = tmp_path / ".tfl_home"
+    ws = tmp_path / ".tfl_ws"
+    home.mkdir()
+    ws.mkdir()
+    monkeypatch.setenv("TFL_HOME_DIR", str(home))
+    monkeypatch.setenv("TFL_WORKSPACE_DIR", str(ws))
+
+    from lab.lab_facade import Lab
+
+    lab = Lab()
+    try:
+        lab.get_dataset("nonexistent_dataset")
+        assert False, "Should have raised FileNotFoundError"
+    except FileNotFoundError:
+        pass  # Expected
+
+
 def test_lab_list_models(tmp_path, monkeypatch):
     _fresh(monkeypatch)
     home = tmp_path / ".tfl_home"