Artifex/control_plane/resources/management/commands/seed_spark_resources.py
2026-08-16 15:35:55 +07:00

109 lines
5 KiB
Python

from __future__ import annotations
import os
import shutil
from pathlib import Path
from django.core.management.base import BaseCommand
from control_plane.resources.models import Resource, ResourceKind
class Command(BaseCommand):
help = "Seed DGX Spark compute and model resources for Sol and Qwen."
def opencode_command(self, model_key: str) -> str:
env_name = f"ARTIFEX_{model_key.upper()}_OPENCODE_COMMAND"
if os.environ.get(env_name):
return os.environ[env_name]
if os.environ.get("ARTIFEX_OPENCODE_COMMAND"):
return os.environ["ARTIFEX_OPENCODE_COMMAND"]
executable = shutil.which("opencode")
if executable is None:
default_user_install = Path.home() / ".opencode" / "bin" / "opencode"
executable = str(default_user_install) if default_user_install.exists() else "opencode"
return f"{executable} run"
def opencode_config(self, model_key: str) -> dict[str, object]:
return {
"model_key": model_key,
"transport": os.environ.get(f"ARTIFEX_{model_key.upper()}_TRANSPORT", os.environ.get("ARTIFEX_OPENCODE_TRANSPORT", "local")),
"command": self.opencode_command(model_key),
"use_pty": os.environ.get(f"ARTIFEX_{model_key.upper()}_USE_PTY", os.environ.get("ARTIFEX_OPENCODE_USE_PTY", "0")) == "1",
"timeout_seconds": int(os.environ.get(f"ARTIFEX_{model_key.upper()}_TIMEOUT_SECONDS", os.environ.get("ARTIFEX_OPENCODE_TIMEOUT_SECONDS", "120"))),
}
def update_opencode_resource(self, *, model_key: str, name: str, roles: list[str], compute: Resource) -> None:
defaults = {
"name": name,
"kind": ResourceKind.MODEL,
"provider": "opencode",
"compute": compute,
"roles": roles,
"config": self.opencode_config(model_key),
}
existing = next((resource for resource in Resource.objects.filter(provider="opencode") if resource.config.get("model_key") == model_key), None)
if existing is not None:
for field, value in defaults.items():
setattr(existing, field, value)
existing.save(update_fields=[*defaults.keys(), "updated_at"])
return
Resource.objects.update_or_create(name=name, defaults={key: value for key, value in defaults.items() if key != "name"})
def handle(self, *args, **options):
spark, _ = Resource.objects.update_or_create(
name=os.environ.get("ARTIFEX_SPARK_RESOURCE_NAME", "DGX Spark"),
defaults={
"kind": ResourceKind.MODEL_COMPUTE,
"provider": "ssh",
"config": {"ssh_alias": os.environ.get("ARTIFEX_SPARK_SSH_ALIAS", "spark")},
"roles": [],
},
)
self.update_opencode_resource(
model_key="sol",
name=os.environ.get("ARTIFEX_SOL_RESOURCE_NAME", "GPT-5.6 Sol"),
roles=["PROJECT_BRAIN", "PLANNING", "ARCHAEOLOGY_INTERPRETATION", "AGENT_DESIGN", "ESCALATION"],
compute=spark,
)
self.update_opencode_resource(
model_key="terra",
name=os.environ.get("ARTIFEX_TERRA_RESOURCE_NAME", "GPT-5.6 Terra"),
roles=["REASONING", "REVIEW", "PORTFOLIO_IC", "STRATEGY"],
compute=spark,
)
self.update_opencode_resource(
model_key="luna",
name=os.environ.get("ARTIFEX_LUNA_RESOURCE_NAME", "GPT-5.6 Luna"),
roles=["REASONING", "MARKET_RESEARCH", "WEB_RESEARCH", "SYNTHESIS"],
compute=spark,
)
Resource.objects.update_or_create(
name=os.environ.get("ARTIFEX_QWEN_RESOURCE_NAME", "Qwen"),
defaults={
"kind": ResourceKind.MODEL,
"provider": "local_inference",
"compute": spark,
"roles": ["CODING", "REVIEW", "REASONING"],
"config": {
"endpoint_url": os.environ.get("ARTIFEX_QWEN_ENDPOINT_URL", "http://192.168.1.162:8002/v1/chat/completions"),
"health_url": os.environ.get("ARTIFEX_QWEN_HEALTH_URL", "http://192.168.1.162:8002/health"),
"model": os.environ.get("ARTIFEX_QWEN_MODEL", "qwen38"),
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
},
},
)
Resource.objects.update_or_create(
name=os.environ.get("ARTIFEX_SEARXNG_RESOURCE_NAME", "SearXNG Search"),
defaults={
"kind": ResourceKind.COMPUTE,
"provider": "searxng",
"compute": spark,
"roles": ["WEB_RESEARCH", "MARKET_RESEARCH", "SEARCH"],
"config": {
"endpoint_url": os.environ.get("ARTIFEX_SEARXNG_ENDPOINT_URL", "http://127.0.0.1:8080"),
"timeout_seconds": int(os.environ.get("ARTIFEX_SEARXNG_TIMEOUT_SECONDS", "20")),
},
},
)
self.stdout.write(self.style.SUCCESS("Seeded Spark model resources."))