Repair TransformerLab activation provisioning

scripts/ensure_transformerlab_user.py

@@ -4,9 +4,15 @@ from __future__ import annotations
 import argparse
 import asyncio
 import os
+import shutil
 import sys
 from pathlib import Path
 
+DEFAULT_WORKSPACE_PLUGINS = ("fastchat_server",)
+DEFAULT_WORKSPACE_EXPERIMENTS = ("alpha", "beta", "gamma")
+DEFAULT_WORKSPACE_MODELS = ("unsloth_Llama-3.2-1B-Instruct",)
+DEFAULT_MODEL_METADATA_FILES = ("_tlab_complete_provenance.json",)
+
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
@@ -36,6 +42,92 @@ def bootstrap_source(transformerlab_dir: Path) -> None:
             os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
 
 
+def target_workspace(transformerlab_dir: Path, team_id: str) -> Path:
+    return transformerlab_dir / "orgs" / team_id / "workspace"
+
+
+def workspace_team_id(workspace: Path, transformerlab_dir: Path) -> str | None:
+    orgs_dir = transformerlab_dir / "orgs"
+    try:
+        relative = workspace.relative_to(orgs_dir)
+    except ValueError:
+        return None
+
+    if len(relative.parts) >= 2 and relative.parts[1] == "workspace":
+        return relative.parts[0]
+    return None
+
+
+def candidate_workspaces(transformerlab_dir: Path, excluded_team_id: str) -> list[Path]:
+    candidates: list[Path] = []
+    root_workspace = transformerlab_dir / "workspace"
+    if root_workspace.is_dir():
+        candidates.append(root_workspace)
+
+    orgs_dir = transformerlab_dir / "orgs"
+    if not orgs_dir.is_dir():
+        return candidates
+
+    for workspace in sorted(orgs_dir.glob("*/workspace")):
+        if not workspace.is_dir():
+            continue
+        if workspace_team_id(workspace, transformerlab_dir) == excluded_team_id:
+            continue
+        candidates.append(workspace)
+    return candidates
+
+
+def copy_dir_if_missing(source: Path | None, target: Path, label: str) -> bool:
+    if source is None or not source.is_dir() or target.exists():
+        return False
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copytree(source, target)
+    print(f"Seeded {label} from {source}.")
+    return True
+
+
+def copy_file_if_missing(source: Path | None, target: Path, label: str) -> bool:
+    if source is None or not source.is_file() or target.exists():
+        return False
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(source, target)
+    print(f"Seeded {label} from {source}.")
+    return True
+
+
+def find_workspace_seed(transformerlab_dir: Path, category: str, name: str, excluded_team_id: str) -> Path | None:
+    for workspace in candidate_workspaces(transformerlab_dir, excluded_team_id):
+        candidate = workspace / category / name
+        if candidate.exists():
+            return candidate
+    return None
+
+
+def seed_workspace(transformerlab_dir: Path, team_id: str) -> None:
+    workspace = target_workspace(transformerlab_dir, team_id)
+    workspace.mkdir(parents=True, exist_ok=True)
+
+    for plugin in DEFAULT_WORKSPACE_PLUGINS:
+        source = transformerlab_dir / "src" / "transformerlab" / "plugins" / plugin
+        copy_dir_if_missing(source, workspace / "plugins" / plugin, f"plugin '{plugin}'")
+
+    for experiment in DEFAULT_WORKSPACE_EXPERIMENTS:
+        source = find_workspace_seed(transformerlab_dir, "experiments", experiment, team_id)
+        copy_dir_if_missing(source, workspace / "experiments" / experiment, f"experiment '{experiment}'")
+
+    copied_model = False
+    for model in DEFAULT_WORKSPACE_MODELS:
+        source = find_workspace_seed(transformerlab_dir, "models", model, team_id)
+        copied_model = copy_dir_if_missing(source, workspace / "models" / model, f"model '{model}'") or copied_model
+
+    for metadata_name in DEFAULT_MODEL_METADATA_FILES:
+        source = find_workspace_seed(transformerlab_dir, "models", metadata_name, team_id)
+        if copied_model or source is not None:
+            copy_file_if_missing(source, workspace / "models" / metadata_name, f"model metadata '{metadata_name}'")
+
+
 async def ensure_user(args: argparse.Namespace) -> int:
     from sqlalchemy import select
     from transformerlab.db.constants import DATABASE_FILE_NAME
@@ -129,6 +221,8 @@ async def ensure_user(args: argparse.Namespace) -> int:
             await session.commit()
             print(f"Updated team role to owner for {args.email}.")
 
+        seed_workspace(Path(args.transformerlab_dir), str(user_team.team_id))
+
         action = "Created" if created else "Verified"
         print(f"{action} default TransformerLab user {args.email}.")
         return 0