Add/model dataset group

api/alembic/versions/a3d2e5f8c901_create_asset_versions_table.py

@@ -0,0 +1,63 @@
+"""create_asset_versions_table
+
+Revision ID: a3d2e5f8c901
+Revises: a1b2c3d4e5f6
+Create Date: 2026-03-06 12:00:00.000000
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "a3d2e5f8c901"
+down_revision: Union[str, Sequence[str], None] = "a1b2c3d4e5f6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Create asset_versions table for tracking versioned groups of models and datasets."""
+    connection = op.get_bind()
+
+    # Helper function to check if table exists
+    def table_exists(table_name: str) -> bool:
+        result = connection.execute(
+            sa.text("SELECT name FROM sqlite_master WHERE type='table' AND name=:name"), {"name": table_name}
+        )
+        return result.fetchone() is not None
+
+    if not table_exists("asset_versions"):
+        op.create_table(
+            "asset_versions",
+            sa.Column("id", sa.String(), nullable=False),
+            sa.Column("asset_type", sa.String(), nullable=False),
+            sa.Column("group_name", sa.String(), nullable=False),
+            sa.Column("version", sa.Integer(), nullable=False),
+            sa.Column("asset_id", sa.String(), nullable=False),
+            sa.Column("tag", sa.String(), nullable=True),
+            sa.Column("job_id", sa.String(), nullable=True),
+            sa.Column("description", sa.String(), nullable=True),
+            sa.Column("created_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
+            sa.PrimaryKeyConstraint("id"),
+        )
+        op.create_index("idx_asset_versions_group", "asset_versions", ["asset_type", "group_name"], unique=False)
+        op.create_index("idx_asset_versions_tag", "asset_versions", ["asset_type", "group_name", "tag"], unique=False)
+        op.create_index("idx_asset_versions_asset_id", "asset_versions", ["asset_id"], unique=False)
+        op.create_index(op.f("ix_asset_versions_asset_type"), "asset_versions", ["asset_type"], unique=False)
+        op.create_index(op.f("ix_asset_versions_group_name"), "asset_versions", ["group_name"], unique=False)
+        op.create_index(op.f("ix_asset_versions_tag_col"), "asset_versions", ["tag"], unique=False)
+
+
+def downgrade() -> None:
+    """Drop asset_versions table."""
+    op.drop_index(op.f("ix_asset_versions_tag_col"), table_name="asset_versions")
+    op.drop_index(op.f("ix_asset_versions_group_name"), table_name="asset_versions")
+    op.drop_index(op.f("ix_asset_versions_asset_type"), table_name="asset_versions")
+    op.drop_index("idx_asset_versions_asset_id", table_name="asset_versions")
+    op.drop_index("idx_asset_versions_tag", table_name="asset_versions")
+    op.drop_index("idx_asset_versions_group", table_name="asset_versions")
+    op.drop_table("asset_versions")

api/api.py

@@ -88,6 +88,7 @@ def _enable_datadog_if_setup():
     api_keys,
     quota,
     ssh_keys,
+    asset_versions,
 )
 from transformerlab.routers.auth import get_user_and_team  # noqa: E402
 
@@ -332,6 +333,7 @@ async def validation_exception_handler(request, exc):
 app.include_router(api_keys.router)
 app.include_router(quota.router)
 app.include_router(ssh_keys.router, dependencies=[Depends(get_user_and_team)])
+app.include_router(asset_versions.router, dependencies=[Depends(get_user_and_team)])
 
 worker_process = None
 

api/transformerlab/routers/asset_versions.py

@@ -0,0 +1,179 @@
+"""
+asset_versions.py
+
+API router for managing versioned groups of models and datasets.
+"""
+
+from typing import Optional
+
+from fastapi import APIRouter, HTTPException, Query
+from pydantic import BaseModel
+
+from transformerlab.services import asset_version_service
+
+
+router = APIRouter(prefix="/asset_versions", tags=["asset_versions"])
+
+
+# ─── Request / Response schemas ───────────────────────────────────────────────
+
+
+class CreateVersionRequest(BaseModel):
+    asset_type: str  # 'model' or 'dataset'
+    group_name: str
+    asset_id: str
+    job_id: Optional[str] = None
+    description: Optional[str] = None
+    tag: Optional[str] = "latest"
+
+
+class SetTagRequest(BaseModel):
+    tag: str  # 'latest', 'production', 'draft'
+
+
+# ─── Group endpoints ─────────────────────────────────────────────────────────
+
+
+@router.get("/groups", summary="List all version groups for a given asset type.")
+async def list_groups(asset_type: str = Query(..., description="'model' or 'dataset'")):
+    try:
+        return await asset_version_service.list_groups(asset_type)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.delete(
+    "/groups/{asset_type}/{group_name}",
+    summary="Delete all versions in a group.",
+)
+async def delete_group(asset_type: str, group_name: str):
+    try:
+        count = await asset_version_service.delete_group(asset_type, group_name)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return {"status": "success", "deleted_count": count}
+
+
+# ─── Version CRUD ─────────────────────────────────────────────────────────────
+
+
+@router.post("/versions", summary="Create a new version in a group.")
+async def create_version(body: CreateVersionRequest):
+    try:
+        result = await asset_version_service.create_version(
+            asset_type=body.asset_type,
+            group_name=body.group_name,
+            asset_id=body.asset_id,
+            job_id=body.job_id,
+            description=body.description,
+            tag=body.tag,
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    return result
+
+
+@router.get(
+    "/versions/{asset_type}/{group_name}",
+    summary="List all versions in a group.",
+)
+async def list_versions(asset_type: str, group_name: str):
+    try:
+        return await asset_version_service.list_versions(asset_type, group_name)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.get(
+    "/versions/{asset_type}/{group_name}/{version}",
+    summary="Get a specific version by number.",
+)
+async def get_version(asset_type: str, group_name: str, version: int):
+    try:
+        result = await asset_version_service.get_version(asset_type, group_name, version)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if result is None:
+        raise HTTPException(status_code=404, detail="Version not found")
+    return result
+
+
+@router.delete(
+    "/versions/{asset_type}/{group_name}/{version}",
+    summary="Delete a specific version.",
+)
+async def delete_version(asset_type: str, group_name: str, version: int):
+    try:
+        deleted = await asset_version_service.delete_version(asset_type, group_name, version)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if not deleted:
+        raise HTTPException(status_code=404, detail="Version not found")
+    return {"status": "success"}
+
+
+# ─── Tag management ──────────────────────────────────────────────────────────
+
+
+@router.put(
+    "/versions/{asset_type}/{group_name}/{version}/tag",
+    summary="Set a tag on a specific version. Moves the tag from any other version in the group.",
+)
+async def set_tag(asset_type: str, group_name: str, version: int, body: SetTagRequest):
+    try:
+        result = await asset_version_service.set_tag(asset_type, group_name, version, body.tag)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if result is None:
+        raise HTTPException(status_code=404, detail="Version not found")
+    return result
+
+
+@router.delete(
+    "/versions/{asset_type}/{group_name}/{version}/tag",
+    summary="Clear the tag from a specific version.",
+)
+async def clear_tag(asset_type: str, group_name: str, version: int):
+    try:
+        result = await asset_version_service.clear_tag(asset_type, group_name, version)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if result is None:
+        raise HTTPException(status_code=404, detail="Version not found")
+    return result
+
+
+# ─── Resolution ──────────────────────────────────────────────────────────────
+
+
+@router.get(
+    "/resolve/{asset_type}/{group_name}",
+    summary="Resolve a group to a specific version. Defaults to 'latest' tag.",
+)
+async def resolve(
+    asset_type: str,
+    group_name: str,
+    tag: Optional[str] = Query(None, description="Tag to resolve: 'latest', 'production', 'draft'"),
+    version: Optional[int] = Query(None, description="Exact version number to resolve"),
+):
+    try:
+        result = await asset_version_service.resolve(asset_type, group_name, tag=tag, version=version)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    if result is None:
+        raise HTTPException(status_code=404, detail="No matching version found")
+    return result
+
+
+# ─── Bulk lookups (used by list views) ────────────────────────────────────────
+
+
+@router.get(
+    "/map/{asset_type}",
+    summary="Get a map of asset_id -> group memberships for annotating list views.",
+)
+async def get_asset_group_map(asset_type: str):
+    try:
+        return await asset_version_service.get_all_asset_group_map(asset_type)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))

api/transformerlab/routers/data.py

@@ -841,6 +841,20 @@ async def dataset_list(generated: bool = True):
     except Exception:
         merged_list = []
 
+    # Augment each dataset with version group info if any
+    try:
+        from transformerlab.services import asset_version_service
+
+        group_map = await asset_version_service.get_all_asset_group_map("dataset")
+        for entry in merged_list:
+            dataset_id = entry.get("dataset_id", "")
+            if dataset_id in group_map:
+                entry["version_groups"] = group_map[dataset_id]
+            else:
+                entry["version_groups"] = []
+    except Exception as e:
+        print(f"Warning: could not fetch dataset version groups: {e}")
+
     if generated:
         return merged_list
 

api/transformerlab/routers/experiment/jobs.py

@@ -35,6 +35,8 @@
     get_job_models_dir,
     get_models_dir,
 )
+from transformerlab.services import asset_version_service
+
 from transformerlab.services.cache_service import cache, cached
 
 router = APIRouter(prefix="/jobs", tags=["train"])
@@ -1103,7 +1105,7 @@ async def get_artifacts(job_id: str, request: Request):
         from lab.dirs import get_job_artifacts_dir
 
         artifacts_dir = await get_job_artifacts_dir(job_id)
-        artifacts = await get_artifacts_from_directory(artifacts_dir, storage)
+        artifacts = await get_artifacts_from_directory(artifacts_dir)
     except Exception as e:
         print(f"Error getting artifacts for job {job_id}: {e}")
         artifacts = []
@@ -1419,7 +1421,11 @@ async def save_dataset_to_registry(
       If a dataset with that name already exists, a timestamped suffix is added.
     - mode='existing': Merge into an existing dataset in the registry. target_name must be provided and must
       refer to an existing dataset. Files from the job dataset are copied into the existing dataset directory.
+
+    In both modes a new version entry is recorded in the asset_versions table
+    so the asset can be tracked as part of a versioned group.
     """
+    from transformerlab.services import asset_version_service
 
     try:
         # Secure the source dataset name
@@ -1471,7 +1477,21 @@ async def save_dataset_to_registry(
             except Exception as copy_err:
                 print(f"Storage.copy_dir failed: {copy_err}")
 
-            return {"status": "success", "message": f"Dataset saved to registry as '{final_name}'"}
+        # Create a version entry for the dataset
+        group_name = dataset_name_secure
+        version_entry = await asset_version_service.create_version(
+            asset_type="dataset",
+            group_name=group_name,
+            asset_id=final_name,
+            job_id=job_id,
+            description=f"Created from job {job_id}",
+        )
+
+        return {
+            "status": "success",
+            "message": f"Dataset saved to registry as '{final_name}'",
+            "version": version_entry,
+        }
 
     except HTTPException:
         raise
@@ -1548,7 +1568,21 @@ async def save_model_to_registry(
             except Exception as copy_err:
                 print(f"storage.copy_dir failed: {copy_err}")
 
-            return {"status": "success", "message": f"Model saved to registry as '{final_name}'"}
+        # Create a version entry for the model
+        group_name = model_name_secure
+        version_entry = await asset_version_service.create_version(
+            asset_type="model",
+            group_name=group_name,
+            asset_id=final_name,
+            job_id=job_id,
+            description=f"Created from job {job_id}",
+        )
+
+        return {
+            "status": "success",
+            "message": f"Model saved to registry as '{final_name}'",
+            "version": version_entry,
+        }
 
     except HTTPException:
         raise

api/transformerlab/routers/model.py

@@ -818,7 +818,23 @@ async def get_model_prompt_template(model: str):
 @router.get("/model/list")
 async def model_local_list(embedding=False):
     # the model list is a combination of downloaded hugging face models and locally generated models
-    return await model_helper.list_installed_models(embedding)
+    models = await model_helper.list_installed_models(embedding)
+
+    # Augment each model with version group info if any
+    try:
+        from transformerlab.services import asset_version_service
+
+        group_map = await asset_version_service.get_all_asset_group_map("model")
+        for model in models:
+            model_id = model.get("model_id", "")
+            if model_id in group_map:
+                model["version_groups"] = group_map[model_id]
+            else:
+                model["version_groups"] = []
+    except Exception as e:
+        print(f"Warning: could not fetch model version groups: {e}")
+
+    return models
 
 
 @router.get("/model/provenance/{model_id}")