wisent-compute 0.3.5__tar.gz

wisent_compute-0.3.5/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: wisent-compute
+Version: 0.3.5
+Summary: Job queue and compute management for Wisent GPU workloads
+Requires-Python: >=3.10
+Requires-Dist: google-cloud-storage>=2.18.0
+Requires-Dist: google-cloud-compute>=1.19.0
+Requires-Dist: google-cloud-pubsub>=2.21.0
+Requires-Dist: google-cloud-secret-manager>=2.20.0
+Requires-Dist: click>=8.0
+Provides-Extra: aws
+Requires-Dist: boto3>=1.34.0; extra == "aws"
+Provides-Extra: dev
+Requires-Dist: functions-framework>=3.0.0; extra == "dev"

wisent_compute-0.3.5/README.md

@@ -0,0 +1,2 @@
+# wisent-compute
+Job queue and compute management for Wisent GPU workloads

wisent_compute-0.3.5/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=68.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "wisent-compute"
+version = "0.3.5"
+description = "Job queue and compute management for Wisent GPU workloads"
+requires-python = ">=3.10"
+dependencies = [
+    "google-cloud-storage>=2.18.0",
+    "google-cloud-compute>=1.19.0",
+    "google-cloud-pubsub>=2.21.0",
+    "google-cloud-secret-manager>=2.20.0",
+    "click>=8.0",
+]
+
+[project.optional-dependencies]
+aws = ["boto3>=1.34.0"]
+dev = ["functions-framework>=3.0.0"]
+
+[project.scripts]
+wc = "wisent_compute.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["wisent_compute*"]
+
+[tool.setuptools.package-data]
+wisent_compute = ["templates/*.sh", "targets/*.json"]

wisent_compute-0.3.5/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

wisent_compute-0.3.5/wisent_compute/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.3.0"

wisent_compute-0.3.5/wisent_compute/cli.py

@@ -0,0 +1,267 @@
+"""CLI entry point: wc submit, wc status, wc results, wc cancel, wc agent."""
+from __future__ import annotations
+
+import json
+import os
+import time
+import urllib.request
+import urllib.error
+
+import click
+
+from .config import (
+    BUCKET,
+    DEFAULT_ANY_PROVIDER,
+    DEFAULT_MAX_COST_PER_HOUR_USD,
+    DEFAULT_PREEMPTIBLE,
+    DEFAULT_PRIORITY,
+)
+from .queue.submit import submit_job, COMPUTE_API
+from .queue.storage import JobStorage
+
+
+def _api_key():
+    return os.environ.get("COMPUTE_API_KEY", "").strip()
+
+
+def _api_get(path):
+    req = urllib.request.Request(
+        f"{COMPUTE_API}{path}",
+        headers={"X-API-Key": _api_key()},
+    )
+    resp = urllib.request.urlopen(req)
+    return json.loads(resp.read())
+
+
+@click.group()
+def main():
+    """Wisent Compute — GPU job queue management."""
+
+
+@main.command()
+@click.argument("command")
+@click.option("--provider", default="gcp",
+              help="Preferred provider (gcp/local). With --any-provider this is just a hint.")
+@click.option("--batch", "batch_file", default=None, help="File with commands")
+@click.option("--spot/--no-spot", default=DEFAULT_PREEMPTIBLE,
+              help="Dispatch on Spot/Preemptible GPUs (cheaper, can be preempted).")
+@click.option("--max-cost-per-hour", "max_cost_per_hour", type=float,
+              default=DEFAULT_MAX_COST_PER_HOUR_USD,
+              help="Hard cap on $/hour for the chosen accelerator. 0 = no cap.")
+@click.option("--any-provider/--pin-provider", "any_provider",
+              default=DEFAULT_ANY_PROVIDER,
+              help="If true (default), any consumer with capacity can claim. "
+                   "If --pin-provider, only the named --provider is allowed.")
+@click.option("--priority", type=int, default=DEFAULT_PRIORITY,
+              help="Higher = scheduled first within FIFO bucket.")
+def submit(command, provider, batch_file, spot, max_cost_per_hour, any_provider, priority):
+    """Submit a job (or batch) to the queue."""
+    commands = []
+    if batch_file:
+        with open(batch_file) as f:
+            commands = [l.strip() for l in f if l.strip() and not l.startswith("#")]
+    else:
+        commands = [command]
+
+    batch_id = f"batch-{int(time.time())}"
+    for cmd in commands:
+        job = submit_job(
+            cmd, provider=provider, batch_id=batch_id, bucket=BUCKET,
+            preemptible=spot,
+            max_cost_per_hour_usd=max_cost_per_hour,
+            pin_to_provider=not any_provider,
+            priority=priority,
+        )
+        click.echo(f"  {job.job_id}  {job.gpu_type or 'cpu':>20s}  {cmd[:60]}")
+    mode = "API" if _api_key() else "GCS"
+    flags = []
+    if spot: flags.append("spot")
+    if max_cost_per_hour > 0: flags.append(f"cap=${max_cost_per_hour:.2f}/hr")
+    if not any_provider: flags.append(f"pinned={provider}")
+    if priority: flags.append(f"priority={priority}")
+    flag_str = (" [" + ", ".join(flags) + "]") if flags else ""
+    click.echo(f"\nSubmitted {len(commands)} job(s) via {mode}{flag_str}. Batch: {batch_id}")
+
+
+@main.command()
+@click.argument("filter_id", required=False)
+def status(filter_id):
+    """Show job status."""
+    if _api_key():
+        _status_api(filter_id)
+    else:
+        _status_gcs(filter_id)
+
+
+def _status_api(filter_id):
+    instances = _api_get("/api/v1/instances")
+    click.echo(f"{'ID':<38} {'STATUS':<12} {'IMAGE':<30} {'COST'}")
+    click.echo("-" * 95)
+    for inst in instances:
+        iid = inst.get("id", "")[:36]
+        st = inst.get("status", "")
+        img = inst.get("docker_image", "")[:28]
+        cost = inst.get("total_cost_cents", 0) / 100
+        if filter_id and filter_id not in iid:
+            continue
+        click.echo(f"{iid:<38} {st:<12} {img:<30} ${cost:.2f}")
+    click.echo(f"\n{len(instances)} instance(s)")
+
+
+def _status_gcs(filter_id):
+    store = JobStorage(BUCKET)
+    all_jobs = store.list_all_jobs()
+    click.echo(f"{'JOB ID':<12} {'STATE':<12} {'GPU':<20} {'RESTARTS':<10} {'COMMAND'}")
+    click.echo("-" * 90)
+    for state in ("running", "queue", "completed", "failed"):
+        for job in all_jobs[state]:
+            if filter_id and filter_id not in (job.job_id, job.batch_id):
+                continue
+            cmd = job.command[:50] + "..." if len(job.command) > 50 else job.command
+            click.echo(f"{job.job_id:<12} {state:<12} {job.gpu_type or 'cpu':<20} {job.restarts:<10} {cmd}")
+    counts = {k: len(v) for k, v in all_jobs.items()}
+    click.echo(f"\n{counts['running']} running, {counts['queue']} queued, "
+               f"{counts['completed']} completed, {counts['failed']} failed")
+
+
+@main.command()
+@click.argument("job_id")
+@click.argument("output_dir")
+def results(job_id, output_dir):
+    """Download job results."""
+    os.makedirs(output_dir, exist_ok=True)
+    os.system(f"gsutil -m cp -r 'gs://{BUCKET}/status/{job_id}/output/*' '{output_dir}/'")
+    click.echo(f"Results downloaded to {output_dir}")
+
+
+@main.command()
+@click.argument("job_id")
+def cancel(job_id):
+    """Cancel a queued or running job."""
+    if _api_key():
+        req = urllib.request.Request(
+            f"{COMPUTE_API}/api/v1/instances/{job_id}",
+            headers={"X-API-Key": _api_key()},
+            method="DELETE",
+        )
+        try:
+            urllib.request.urlopen(req)
+            click.echo(f"Cancelled {job_id}")
+        except urllib.error.HTTPError as e:
+            click.echo(f"Failed: {e.code}")
+        return
+
+    store = JobStorage(BUCKET)
+    job = store.read_job("queue", job_id)
+    if job:
+        store.delete_job("queue", job_id)
+        click.echo(f"Removed {job_id} from queue")
+        return
+    job = store.read_job("running", job_id)
+    if job and job.instance_ref:
+        from .providers import get_provider
+        prov = get_provider(job.provider)
+        prov.delete_instance(job.instance_ref)
+        job.state = "failed"
+        job.error = "cancelled"
+        store.move_job(job, "running", "failed")
+        click.echo(f"Cancelled {job_id}, instance terminated")
+        return
+    click.echo(f"Job {job_id} not found")
+
+
+@main.command()
+@click.option("--gpu-type", default="", help="GPU type (auto-detected if --target/--auto absent)")
+@click.option("--target", default=None,
+              help="Pull gpu_type and slot count from registry by name.")
+@click.option("--auto", is_flag=True, default=False,
+              help="Look up self in the GCS-hosted registry by hostname; no manual config.")
+def agent(gpu_type, target, auto):
+    """Run local GPU agent. Polls queue, respects Vast.ai renters.
+
+    --auto looks up the local hostname in gs://wisent-compute/registry.json
+    and uses that entry's slots/gpu_type. Re-fetches periodically so registry
+    edits propagate without restarting the agent.
+    """
+    import os as _os
+    if auto:
+        from .targets import lookup_self
+        t = lookup_self(_os.uname().nodename, source="gcs")
+        if not t:
+            raise click.ClickException(
+                f"hostname '{_os.uname().nodename}' not found in GCS registry"
+            )
+        gpu_type = gpu_type or (t.gpu_type or "")
+        _os.environ["WC_LOCAL_SLOTS"] = str(t.slots)
+        click.echo(f"agent --auto: target={t.name} gpu_type={gpu_type} slots={t.slots}")
+    elif target:
+        from .targets import lookup
+        t = lookup(target)
+        if not t:
+            raise click.ClickException(f"target '{target}' not found in registry")
+        if t.kind != "local":
+            raise click.ClickException(f"target '{target}' kind={t.kind}, expected local")
+        gpu_type = gpu_type or (t.gpu_type or "")
+        _os.environ["WC_LOCAL_SLOTS"] = str(t.slots)
+        click.echo(f"agent: target={t.name} gpu_type={gpu_type} slots={t.slots}")
+    from .providers.local_agent import run_agent
+    run_agent(gpu_type=gpu_type)
+
+
+@main.command()
+@click.option("--target", default=None,
+              help="Coordinator name in registry (default: the one with active=true).")
+@click.option("--once", is_flag=True, default=False,
+              help="Run a single scheduling tick and exit (cron-friendly).")
+def coordinator(target, once):
+    """Run the scheduling tick locally instead of the GCP Cloud Function.
+
+    Reads the named coordinator entry from the registry, loops on its
+    interval_seconds, runs the same monitor_jobs/schedule_queued_jobs
+    chain the Cloud Function does. State stays in the registry-declared
+    state_uri so all consumers (cloud + local) keep seeing the same queue.
+    """
+    from .coordinator import run as run_coordinator
+    raise SystemExit(run_coordinator(target=target, once=once))
+
+
+@main.group()
+def registry():
+    """Manage the canonical compute-target registry hosted in GCS."""
+
+
+@registry.command("push")
+@click.argument("path", type=click.Path(exists=True, dir_okay=False), required=False)
+def registry_push(path):
+    """Upload local registry.json to gs://wisent-compute/registry.json."""
+    import shutil, subprocess
+    from .targets import REGISTRY_PATH, GCS_REGISTRY_URI
+    src = path or str(REGISTRY_PATH)
+    gsutil = shutil.which("gsutil") or "gsutil"
+    r = subprocess.run([gsutil, "cp", src, GCS_REGISTRY_URI], capture_output=True, text=True)
+    if r.returncode != 0:
+        raise click.ClickException(r.stderr or r.stdout or "gsutil cp failed")
+    click.echo(f"pushed {src} -> {GCS_REGISTRY_URI}")
+
+
+@registry.command("pull")
+def registry_pull():
+    """Print the GCS-hosted registry to stdout."""
+    from .targets import _load_from_gcs
+    data = _load_from_gcs()
+    if data is None:
+        raise click.ClickException("could not fetch registry from GCS")
+    import json as _json
+    click.echo(_json.dumps(data, indent=2))
+
+
+@main.command()
+@click.option("--target", default=None, help="Specific entry name (target or coordinator).")
+@click.option("--dry-run", is_flag=True, default=False, help="Print unit/plist; do not enable.")
+@click.option("--local", "local_install", is_flag=True, default=False,
+              help="Install on THIS machine (launchd/systemd --user) instead of via SSH.")
+def bootstrap(target, dry_run, local_install):
+    """Provision wisent-compute services persistently across reboots."""
+    from .deploy.bootstrap import run_bootstrap
+    run_bootstrap(target=target, dry_run=dry_run, local_install=local_install,
+                  echo=click.echo)

wisent_compute-0.3.5/wisent_compute/cloud_function/main.py

@@ -0,0 +1,56 @@
+"""Cloud Function entry point. Triggered every 3 min by Cloud Scheduler."""
+from __future__ import annotations
+
+import sys
+from google.cloud import pubsub_v1, secretmanager_v1
+
+from wisent_compute.config import PROJECT, BUCKET, ALERTS_TOPIC
+from wisent_compute.queue.storage import JobStorage
+from wisent_compute.providers import get_provider
+from wisent_compute.monitor import check_running_jobs
+from wisent_compute.scheduler import schedule_queued_jobs
+
+_publisher = None
+_secrets = None
+
+
+def _log(msg):
+    sys.stderr.write(f"[tick] {msg}\n")
+    sys.stderr.flush()
+
+
+def _load_secrets():
+    global _secrets
+    if _secrets is not None:
+        return _secrets
+    client = secretmanager_v1.SecretManagerServiceClient()
+    _secrets = {}
+    for name in ("wisent-hf-token", "wisent-gh-token"):
+        try:
+            r = client.access_secret_version(request={
+                "name": f"projects/{PROJECT}/secrets/{name}/versions/latest"
+            })
+            key = name.replace("wisent-", "").replace("-", "_").upper()
+            _secrets[key] = r.payload.data.decode("utf-8")
+        except Exception:
+            pass
+    return _secrets
+
+
+def monitor_jobs(request=None):
+    """Main tick: check running, then schedule queued."""
+    global _publisher
+    _log("Tick started")
+
+    store = JobStorage(BUCKET)
+    provider = get_provider("gcp")
+    if _publisher is None:
+        _publisher = pubsub_v1.PublisherClient()
+
+    check_running_jobs(store, provider, _publisher)
+
+    secrets = _load_secrets()
+    scheduled = schedule_queued_jobs(store, provider, "gcp", secrets)
+    _log(f"Tick done: scheduled {scheduled}")
+
+    return "OK"

wisent_compute-0.3.5/wisent_compute/config.py

@@ -0,0 +1,90 @@
+"""Configuration and constants."""
+from __future__ import annotations
+
+import os
+import json
+import re
+
+PROJECT = os.environ.get("GCP_PROJECT", "wisent-480400")
+BUCKET = os.environ.get("WC_BUCKET", "wisent-compute")
+REGION = os.environ.get("GCP_REGION", "us-central1")
+ALERTS_TOPIC = os.environ.get("WC_ALERTS_TOPIC", f"projects/{PROJECT}/topics/wisent-compute-alerts")
+
+ZONE_ROTATION = [f"{REGION}-b", f"{REGION}-a", f"{REGION}-c", f"{REGION}-f"]
+HEARTBEAT_STALE_MINUTES = 15
+MAX_SCHEDULE_PER_TICK = 4
+INSTANCE_PREFIX = "wisent"
+
+# Defaults for the smart-routing CLI flags. 0 means "no cap"; the scheduler
+# only enforces a cost gate when this is positive.
+DEFAULT_MAX_COST_PER_HOUR_USD = 0.0
+DEFAULT_PRIORITY = 0
+DEFAULT_PREEMPTIBLE = False
+DEFAULT_ANY_PROVIDER = True
+
+DEFAULT_IMAGE = "pytorch-2-9-cu129-ubuntu-2204-nvidia-580-v20260408"
+DEFAULT_IMAGE_PROJECT = "deeplearning-platform-release"
+DEFAULT_CPU_IMAGE_FAMILY = "ubuntu-2204-lts"
+DEFAULT_CPU_IMAGE_PROJECT = "ubuntu-os-cloud"
+DEFAULT_BOOT_DISK_GB = 200
+
+
+def estimate_gpu_memory(command: str) -> int:
+    """Estimate GPU memory needed from a command string."""
+    model_match = re.search(r'--model\s+(\S+)', command)
+    if not model_match:
+        return 0
+    model = model_match.group(1)
+
+    params_b = 0
+    m = re.search(r'(\d+\.?\d*)[Bb]', model)
+    if m:
+        params_b = float(m.group(1))
+    else:
+        m = re.search(r'(\d+)[Mm]', model)
+        if m:
+            params_b = int(m.group(1)) / 1000
+    if params_b == 0:
+        params_b = 7
+
+    quant_factor = 1
+    if re.search(r'GPTQ|AWQ|INT4|4bit|Q4', model):
+        quant_factor = 4
+    elif re.search(r'INT8|8bit|Q8', model):
+        quant_factor = 2
+
+    weights_gb = params_b * 2 / quant_factor
+    kv_gb = weights_gb * 0.3
+    overhead_gb = 8
+
+    multiplier = 1.0
+    if re.search(r'get-activations|generate-vector', command):
+        multiplier = 1.2
+    elif re.search(r'modify-weights|optimize-weights|training', command):
+        multiplier = 1.5
+
+    return round((weights_gb + kv_gb + overhead_gb) * multiplier)
+
+
+def lookup_instance_type(provider: str, gpu_mem_gb: int) -> tuple[str, str]:
+    """Return (machine_type, accel_type) for the given memory requirement.
+
+    If gpu_mem_gb exceeds every tier in GPU_SIZING, returns the LARGEST
+    available tier rather than ("", ""). The previous behavior produced an
+    empty machine_type that the GCE create_instance call rejected with
+    'Machine type with name "" does not exist', wedging the job. Sending
+    it to the largest tier means the in-VM workload may still OOM, but
+    that's a clearer failure mode than a malformed instance request.
+    """
+    from .models import GPU_SIZING
+    sizing = GPU_SIZING.get(provider, {})
+    if not sizing:
+        return ("", "")
+    best_mem, best_spec = 10**9, None
+    largest_mem, largest_spec = 0, None
+    for mem, spec in sizing.items():
+        if mem > largest_mem:
+            largest_mem, largest_spec = mem, spec
+        if mem >= gpu_mem_gb and mem < best_mem:
+            best_mem, best_spec = mem, spec
+    return best_spec if best_spec is not None else largest_spec

wisent_compute-0.3.5/wisent_compute/coordinator.py

@@ -0,0 +1,93 @@
+"""Coordinator daemon - the same scheduling tick the GCP Cloud Function runs,
+as a long-lived local process so the system can run without GCP CF / Cloud
+Scheduler.
+
+Reads the named coordinator entry from the registry (default: the one whose
+active=true), constructs the same JobStorage + GCP provider + scheduler call
+chain monitor_jobs uses, and loops on the configured interval_seconds.
+
+--once runs a single tick and exits (cron-driven runtimes).
+
+State stays in the registry-declared state_uri (currently always GCS), so
+swapping coordinator from GCF to a daemon on the Mac doesn't change which
+queue the agents see.
+"""
+from __future__ import annotations
+
+import sys
+import time
+from typing import Optional
+
+from .config import BUCKET
+from .monitor import check_running_jobs
+from .providers import get_provider
+from .queue.storage import JobStorage
+from .scheduler import schedule_queued_jobs
+from .targets import Coordinator, load_coordinators, lookup_coordinator
+
+
+def _log(msg: str) -> None:
+    sys.stderr.write(f"[tick] {msg}\n")
+    sys.stderr.flush()
+
+
+def _resolve_coordinator(target: Optional[str]) -> Coordinator:
+    """Pick the coordinator entry: explicit --target, or the active one."""
+    if target:
+        c = lookup_coordinator(target)
+        if c is None:
+            raise SystemExit(f"coordinator '{target}' not found in registry")
+        return c
+    active = [c for c in load_coordinators() if c.active]
+    if not active:
+        raise SystemExit(
+            "no active coordinator in registry. Set active=true on one entry "
+            "or pass --target NAME explicitly."
+        )
+    if len(active) > 1:
+        names = ", ".join(c.name for c in active)
+        raise SystemExit(f"multiple active coordinators ({names}); set active=true on exactly one")
+    return active[0]
+
+
+def _bucket_from_state_uri(state_uri: str) -> str:
+    """Strip 'gs://' prefix to get the bucket name JobStorage expects."""
+    if state_uri.startswith("gs://"):
+        return state_uri[len("gs://"):].split("/", 1)[0]
+    return state_uri
+
+
+def _run_tick(store: JobStorage, secrets: dict) -> int:
+    """One scheduling cycle. Returns the number of jobs newly scheduled."""
+    provider = get_provider("gcp")
+    check_running_jobs(store, provider, publisher=None)
+    scheduled = schedule_queued_jobs(store, provider, "gcp", secrets)
+    return scheduled
+
+
+def run(target: Optional[str] = None, once: bool = False) -> int:
+    """Coordinator daemon entry point. Used by `wc coordinator`."""
+    coord = _resolve_coordinator(target)
+    if coord.runtime == "gcp_cloud_function":
+        _log(
+            f"coordinator '{coord.name}' runtime=gcp_cloud_function: tick is "
+            f"driven by Cloud Scheduler, this daemon is a no-op. Use --target "
+            f"to point at a runtime=daemon entry instead."
+        )
+        return 0
+
+    bucket = _bucket_from_state_uri(coord.state_uri) or BUCKET
+    store = JobStorage(bucket)
+    interval = max(15, int(coord.interval_seconds))
+    _log(f"coordinator '{coord.name}' runtime={coord.runtime} interval={interval}s state={coord.state_uri}")
+
+    secrets: dict = {}
+    while True:
+        try:
+            n = _run_tick(store, secrets)
+            _log(f"tick scheduled={n}")
+        except Exception as exc:
+            _log(f"tick failed: {exc!r}")
+        if once:
+            return 0
+        time.sleep(interval)