|
| 1 | +import shutil |
| 2 | +import threading |
| 3 | +from pathlib import Path |
| 4 | +from typing import List, Optional |
| 5 | + |
| 6 | +import git |
| 7 | +import yaml |
| 8 | +from cachetools import TTLCache, cached |
| 9 | + |
| 10 | +from dstack._internal.core.models.templates import UITemplate |
| 11 | +from dstack._internal.server import settings |
| 12 | +from dstack._internal.utils.common import run_async |
| 13 | +from dstack._internal.utils.logging import get_logger |
| 14 | + |
| 15 | +logger = get_logger(__name__) |
| 16 | + |
| 17 | +TEMPLATES_DIR_NAME = ".dstack/templates" |
| 18 | +CACHE_TTL_SECONDS = 180 |
| 19 | + |
| 20 | +_repo_path: Optional[Path] = None |
| 21 | +_templates_cache: TTLCache = TTLCache(maxsize=1, ttl=CACHE_TTL_SECONDS) |
| 22 | +_templates_lock = threading.Lock() |
| 23 | + |
| 24 | + |
| 25 | +async def list_templates() -> List[UITemplate]: |
| 26 | + """Return templates available for the UI. |
| 27 | +
|
| 28 | + Currently returns only server-wide templates configured via DSTACK_SERVER_TEMPLATES_REPO. |
| 29 | + Project-specific templates will be included once implemented. |
| 30 | + """ |
| 31 | + if not settings.SERVER_TEMPLATES_REPO: |
| 32 | + return [] |
| 33 | + return await run_async(_list_templates_sync) |
| 34 | + |
| 35 | + |
| 36 | +@cached(cache=_templates_cache, lock=_templates_lock) |
| 37 | +def _list_templates_sync() -> List[UITemplate]: |
| 38 | + _fetch_templates_repo() |
| 39 | + return _parse_templates() |
| 40 | + |
| 41 | + |
| 42 | +def _fetch_templates_repo() -> None: |
| 43 | + global _repo_path |
| 44 | + |
| 45 | + repo_dir = settings.SERVER_DATA_DIR_PATH / "templates-repo" |
| 46 | + |
| 47 | + if _repo_path is not None and _repo_path.exists(): |
| 48 | + repo = git.Repo(str(_repo_path)) |
| 49 | + repo.remotes.origin.pull() |
| 50 | + return |
| 51 | + |
| 52 | + if repo_dir.exists(): |
| 53 | + try: |
| 54 | + repo = git.Repo(str(repo_dir)) |
| 55 | + repo.remotes.origin.pull() |
| 56 | + _repo_path = repo_dir |
| 57 | + return |
| 58 | + except (git.InvalidGitRepositoryError, git.GitCommandError): |
| 59 | + logger.warning("Invalid templates repo at %s, re-cloning", repo_dir) |
| 60 | + shutil.rmtree(repo_dir) |
| 61 | + |
| 62 | + assert settings.SERVER_TEMPLATES_REPO is not None |
| 63 | + git.Repo.clone_from( |
| 64 | + settings.SERVER_TEMPLATES_REPO, |
| 65 | + str(repo_dir), |
| 66 | + depth=1, |
| 67 | + ) |
| 68 | + _repo_path = repo_dir |
| 69 | + |
| 70 | + |
| 71 | +def _parse_templates() -> List[UITemplate]: |
| 72 | + if _repo_path is None: |
| 73 | + return [] |
| 74 | + |
| 75 | + templates_dir = _repo_path / TEMPLATES_DIR_NAME |
| 76 | + if not templates_dir.is_dir(): |
| 77 | + logger.warning("Templates directory %s not found in repo", TEMPLATES_DIR_NAME) |
| 78 | + return [] |
| 79 | + |
| 80 | + templates: List[UITemplate] = [] |
| 81 | + for entry in sorted(templates_dir.iterdir()): |
| 82 | + if entry.suffix not in (".yml", ".yaml"): |
| 83 | + continue |
| 84 | + try: |
| 85 | + with open(entry) as f: |
| 86 | + data = yaml.safe_load(f) |
| 87 | + if not isinstance(data, dict): |
| 88 | + logger.warning("Skipping %s: not a valid YAML mapping", entry.name) |
| 89 | + continue |
| 90 | + if data.get("type") != "ui-template": |
| 91 | + logger.debug("Skipping %s: type is not 'ui-template'", entry.name) |
| 92 | + continue |
| 93 | + template = UITemplate.parse_obj(data) |
| 94 | + templates.append(template) |
| 95 | + except Exception: |
| 96 | + logger.warning("Skipping invalid template %s", entry.name, exc_info=True) |
| 97 | + continue |
| 98 | + |
| 99 | + return templates |
0 commit comments