Merge pull request #1194 from transformerlab/add/team-secrets

Add access for team secrets

Author: deep1401

api/transformerlab/routers/compute_provider.py

@@ -2,6 +2,7 @@
 
 import os
 import time
+import json
 import configparser
 from pathlib import Path
 from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File, Request
@@ -27,8 +28,7 @@
     ProviderTemplateFileUploadResponse,
     ResumeFromCheckpointRequest,
 )
-from transformerlab.shared.models.models import ProviderType, TeamComputeProvider
-from transformerlab.compute_providers.base import ComputeProvider
+from transformerlab.shared.models.models import ProviderType
 from transformerlab.compute_providers.models import (
     ClusterConfig,
     ClusterStatus,
@@ -45,6 +45,7 @@
     read_github_pat_from_workspace,
     generate_github_clone_setup,
 )
+from transformerlab.shared.secret_utils import load_team_secrets, replace_secrets_in_dict, replace_secret_placeholders
 from typing import Any
 
 router = APIRouter(prefix="/compute_provider", tags=["compute_provider"])
@@ -59,17 +60,6 @@ def _sanitize_cluster_basename(base_name: Optional[str]) -> str:
     return normalized or "remote-template"
 
 
-def _get_provider_instances(providers: list[TeamComputeProvider]) -> Dict[str, ComputeProvider]:
-    """Instantiate providers safely."""
-    instances: Dict[str, ComputeProvider] = {}
-    for provider in providers:
-        try:
-            instances[provider.id] = get_provider_instance(provider)
-        except Exception as exc:
-            print(f"Failed to instantiate provider {provider.id}: {exc}")
-    return instances
-
-
 @router.post("/{provider_id}/task/{task_id}/file-upload", response_model=ProviderTemplateFileUploadResponse)
 async def upload_task_file_for_provider(
     provider_id: str,
@@ -232,7 +222,6 @@ async def get_usage_report(
     Aggregates usage data by user, provider, and resources.
     Only accessible to team owners.
     """
-    import json
     from datetime import datetime
     from lab import Experiment
 
@@ -562,46 +551,6 @@ async def check_provider(
 # ============================================================================
 
 
-@router.post("/{provider_id}/clusters/{cluster_name}/launch")
-async def launch_cluster(
-    provider_id: str,
-    cluster_name: str,
-    config: ClusterConfig,
-    user_and_team=Depends(get_user_and_team),
-    session: AsyncSession = Depends(get_async_session),
-):
-    """
-    Launch/provision a new cluster using the specified provider.
-    Requires X-Team-Id header and team membership.
-    """
-    team_id = user_and_team["team_id"]
-
-    provider = await get_team_provider(session, team_id, provider_id)
-    if not provider:
-        raise HTTPException(status_code=404, detail="Provider not found")
-
-    try:
-        # Get provider instance
-        provider_instance = get_provider_instance(provider)
-
-        if not config.cluster_name:
-            config.cluster_name = cluster_name
-        config.provider_name = provider.name
-        config.provider_id = provider.id
-
-        # Launch cluster
-        result = provider_instance.launch_cluster(cluster_name, config)
-
-        return {
-            "status": "success",
-            "message": f"Cluster '{cluster_name}' launch initiated",
-            "cluster_name": cluster_name,
-            "result": result,
-        }
-    except Exception:
-        raise HTTPException(status_code=500, detail="Failed to launch cluster")
-
-
 def _get_aws_credentials_from_file(profile_name: str = "transformerlab-s3") -> Tuple[Optional[str], Optional[str]]:
     """
     Read AWS credentials from ~/.aws/credentials file for the specified profile.
@@ -802,6 +751,9 @@ async def _launch_sweep_jobs(
 
             provider_display_name = request.provider_name or provider.name
 
+            # Load team secrets for template replacement
+            team_secrets = await load_team_secrets()
+
             # Generate all parameter combinations
             param_names = list(sweep_config.keys())
             param_values = [sweep_config[name] for name in param_names]
@@ -832,6 +784,11 @@ async def _launch_sweep_jobs(
 
                 # Prepare environment variables
                 env_vars = request.env_vars.copy() if request.env_vars else {}
+
+                # Replace {{secret.<name>}} patterns in env_vars
+                if env_vars and team_secrets:
+                    env_vars = replace_secrets_in_dict(env_vars, team_secrets)
+
                 env_vars["_TFL_JOB_ID"] = str(child_job_id)
                 env_vars["_TFL_EXPERIMENT_ID"] = request.experiment_id
 
@@ -872,11 +829,25 @@ async def _launch_sweep_jobs(
                     )
                     setup_commands.append(github_setup)
 
+                # Add user-provided setup if any (replace secrets in setup)
                 if request.setup:
-                    setup_commands.append(request.setup)
+                    setup_with_secrets = (
+                        replace_secret_placeholders(request.setup, team_secrets) if team_secrets else request.setup
+                    )
+                    setup_commands.append(setup_with_secrets)
 
                 final_setup = ";".join(setup_commands) if setup_commands else None
 
+                # Replace secrets in command
+                command_with_secrets = (
+                    replace_secret_placeholders(request.command, team_secrets) if team_secrets else request.command
+                )
+
+                # Replace secrets in parameters if present
+                parameters_with_secrets = merged_params
+                if merged_params and team_secrets:
+                    parameters_with_secrets = replace_secrets_in_dict(merged_params, team_secrets)
+
                 # Store child job data
                 child_job_data = {
                     "parent_sweep_job_id": str(parent_job_id),
@@ -886,7 +857,7 @@ async def _launch_sweep_jobs(
                     "task_name": f"{request.task_name or 'Task'} (Sweep {i + 1}/{total_configs})"
                     if request.task_name
                     else None,
-                    "command": request.command,
+                    "command": command_with_secrets,
                     "cluster_name": formatted_cluster_name,
                     "subtype": request.subtype,
                     "cpus": request.cpus,
@@ -897,7 +868,7 @@ async def _launch_sweep_jobs(
                     "setup": final_setup,
                     "env_vars": env_vars if env_vars else None,
                     "file_mounts": request.file_mounts or None,
-                    "parameters": merged_params or None,
+                    "parameters": parameters_with_secrets or None,
                     "provider_id": provider.id,
                     "provider_type": provider.type,
                     "provider_name": provider_display_name,
@@ -923,7 +894,7 @@ async def _launch_sweep_jobs(
                     cluster_name=formatted_cluster_name,
                     provider_name=provider_display_name,
                     provider_id=provider.id,
-                    command=request.command,
+                    command=command_with_secrets,
                     setup=final_setup,
                     env_vars=env_vars,
                     cpus=request.cpus,
@@ -1105,9 +1076,16 @@ async def launch_template_on_provider(
 
     provider_display_name = request.provider_name or provider.name
 
+    # Load team secrets for template replacement
+    team_secrets = await load_team_secrets()
+
     # Prepare environment variables - start with a copy of requested env_vars
     env_vars = request.env_vars.copy() if request.env_vars else {}
 
+    # Replace {{secret.<name>}} patterns in env_vars
+    if env_vars and team_secrets:
+        env_vars = replace_secrets_in_dict(env_vars, team_secrets)
+
     # Get AWS credentials from stored credentials file (transformerlab-s3 profile)
     aws_profile = "transformerlab-s3"
     if os.getenv("TFL_API_STORAGE_URI"):
@@ -1132,9 +1110,10 @@ async def launch_template_on_provider(
         )
         setup_commands.append(github_setup)
 
-    # Add user-provided setup if any
+    # Add user-provided setup if any (replace secrets in setup)
     if request.setup:
-        setup_commands.append(request.setup)
+        setup_with_secrets = replace_secret_placeholders(request.setup, team_secrets) if team_secrets else request.setup
+        setup_commands.append(setup_with_secrets)
 
     final_setup = ";".join(setup_commands) if setup_commands else None
 
@@ -1159,9 +1138,21 @@ async def launch_template_on_provider(
         # env_vars["AWS_ACCESS_KEY_ID"] = aws_access_key_id
         # env_vars["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
 
+    # Replace secrets in command
+    command_with_secrets = (
+        replace_secret_placeholders(request.command, team_secrets) if team_secrets else request.command
+    )
+
+    # Replace secrets in parameters if present
+    parameters_with_secrets = None
+    if request.parameters and team_secrets:
+        parameters_with_secrets = replace_secrets_in_dict(request.parameters, team_secrets)
+    else:
+        parameters_with_secrets = request.parameters
+
     job_data = {
         "task_name": request.task_name,
-        "command": request.command,
+        "command": command_with_secrets,
         "cluster_name": formatted_cluster_name,
         "subtype": request.subtype,
         "interactive_type": request.interactive_type,
@@ -1173,7 +1164,7 @@ async def launch_template_on_provider(
         "setup": final_setup,
         "env_vars": env_vars if env_vars else None,
         "file_mounts": request.file_mounts or None,
-        "parameters": request.parameters or None,
+        "parameters": parameters_with_secrets or None,
         "provider_id": provider.id,
         "provider_type": provider.type,
         "provider_name": provider_display_name,
@@ -1197,7 +1188,7 @@ async def launch_template_on_provider(
         cluster_name=formatted_cluster_name,
         provider_name=provider_display_name,
         provider_id=provider.id,
-        command=request.command,
+        command=command_with_secrets,
         setup=final_setup,
         env_vars=env_vars,
         cpus=request.cpus,

api/transformerlab/routers/teams.py

@@ -1,4 +1,4 @@
-from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Query
 from fastapi.responses import FileResponse, Response
 from sqlalchemy.ext.asyncio import AsyncSession
 from transformerlab.utils.api_key_utils import mask_key
@@ -9,13 +9,14 @@
 from transformerlab.utils.email import send_team_invitation_email
 from transformerlab.shared.remote_workspace import create_bucket_for_team
 
-from pydantic import BaseModel, EmailStr
+from pydantic import BaseModel, EmailStr, Field
 from typing import Optional
 from sqlalchemy import select, delete, update, func, and_
 from datetime import datetime, timedelta
 from os import getenv
 from PIL import Image
 import io
+import json
 import logging
 
 from lab import Experiment
@@ -71,6 +72,10 @@ class GitHubPATRequest(BaseModel):
     pat: Optional[str] = None
 
 
+class TeamSecretsRequest(BaseModel):
+    secrets: dict[str, str] = Field(..., description="Team secrets as key-value pairs")
+
+
 router = APIRouter(tags=["teams"])
 
 
@@ -1019,3 +1024,99 @@ async def delete_team_logo(
     except Exception as e:
         print(f"Error deleting team logo: {e}")
         raise HTTPException(status_code=500, detail="Failed to delete team logo")
+
+
+@router.get("/teams/{team_id}/secrets")
+async def get_team_secrets(
+    team_id: str,
+    user_and_team=Depends(get_user_and_team),
+    include_values: bool = Query(False, description="Include actual secret values (only for team owners)"),
+):
+    """
+    Get team secrets.
+    - Team members can view secret keys only (values are masked).
+    - Team owners can view actual values by setting include_values=true.
+    """
+    # Verify team_id matches the one in header
+    if team_id != user_and_team["team_id"]:
+        raise HTTPException(status_code=400, detail="Team ID mismatch")
+
+    workspace_dir = await get_workspace_dir()
+    secrets_path = storage.join(workspace_dir, "team_secrets.json")
+
+    try:
+        if not await storage.exists(secrets_path):
+            return {"status": "success", "secrets": {}}
+
+        async with await storage.open(secrets_path, "r") as f:
+            secrets = json.loads(await f.read())
+
+        # Check if user is team owner and requested values
+        is_owner = user_and_team.get("role") == TeamRole.OWNER.value
+        if include_values and is_owner:
+            # Return actual values for team owners
+            return {
+                "status": "success",
+                "secrets": secrets,
+                "secret_keys": list(secrets.keys()),
+            }
+        else:
+            # Mask all secret values for security
+            masked_secrets = {key: "***" for key in secrets.keys()}
+            return {
+                "status": "success",
+                "secrets": masked_secrets,
+                "secret_keys": list(secrets.keys()),
+            }
+    except Exception as e:
+        print(f"Error reading team secrets: {e}")
+        raise HTTPException(status_code=500, detail="Failed to read team secrets")
+
+
+@router.put("/teams/{team_id}/secrets")
+async def set_team_secrets(
+    team_id: str,
+    secrets_data: TeamSecretsRequest,
+    owner_info=Depends(require_team_owner),
+):
+    """
+    Set team secrets. Only team owners can set/update secrets.
+    Stored in workspace/team_secrets.json file.
+    """
+    # Verify team_id matches the one in header
+    if team_id != owner_info["team_id"]:
+        raise HTTPException(status_code=400, detail="Team ID mismatch")
+
+    workspace_dir = await get_workspace_dir()
+    secrets_path = storage.join(workspace_dir, "team_secrets.json")
+
+    try:
+        # Validate that all keys are valid environment variable names
+        # Environment variable names can contain letters, numbers, and underscores
+        import re
+
+        valid_key_pattern = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+        for key in secrets_data.secrets.keys():
+            if not valid_key_pattern.match(key):
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Invalid secret key '{key}'. Secret keys must start with a letter or underscore and contain only letters, numbers, and underscores.",
+                )
+
+        # Ensure workspace directory exists
+        await storage.makedirs(workspace_dir, exist_ok=True)
+
+        # Write secrets to file
+        async with await storage.open(secrets_path, "w") as f:
+            await f.write(json.dumps(secrets_data.secrets, indent=2))
+
+        return {
+            "status": "success",
+            "message": "Team secrets saved successfully",
+            "secret_keys": list(secrets_data.secrets.keys()),
+        }
+    except HTTPException:
+        raise
+    except Exception as e:
+        print(f"Error saving team secrets: {e}")
+        raise HTTPException(status_code=500, detail="Failed to save team secrets")