relaymd-cli 0.1.0__tar.gz

relaymd_cli-0.1.0/.gitignore

@@ -0,0 +1,30 @@
+# Python
+__pycache__/
+*.py[cod]
+*.so
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+*.egg-info/
+*.egg
+*.whl
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# IDE
+.idea/
+.vscode/
+
+# Environment files
+.env
+.env.*
+!.env.orchestrator.example
+!.env.worker.example
+
+# Local databases
+relaymd.db

relaymd_cli-0.1.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: relaymd-cli
+Version: 0.1.0
+Requires-Python: >=3.11
+Requires-Dist: httpx
+Requires-Dist: pydantic-settings[yaml]
+Requires-Dist: pyyaml
+Requires-Dist: relaymd-models
+Requires-Dist: relaymd-storage
+Requires-Dist: rich
+Requires-Dist: typer

relaymd_cli-0.1.0/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "relaymd-cli"
+version = "0.1.0"
+requires-python = ">=3.11"
+dependencies = [
+    "relaymd-models",
+    "relaymd-storage",
+    "httpx",
+    "typer",
+    "rich",
+    "pydantic-settings[yaml]",
+    "pyyaml",
+]
+
+[project.scripts]
+relaymd = "relaymd.cli.__main__:app"
+
+[tool.uv.sources]
+relaymd-models = { workspace = true }
+relaymd-storage = { workspace = true }
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/relaymd"]

relaymd_cli-0.1.0/relaymd-cli.spec

@@ -0,0 +1,28 @@
+# -*- mode: python ; coding: utf-8 -*-
+a = Analysis(
+    ['src/relaymd/cli/__main__.py'],
+    pathex=['src'],
+    binaries=[],
+    datas=[],
+    hiddenimports=['relaymd.models', 'relaymd.storage'],
+    hookspath=[],
+    hooksconfig={},
+    runtime_hooks=[],
+    excludes=[],
+    noarchive=False,
+)
+pex = PYZ(a.pure)
+exe = EXE(
+    pex,
+    a.scripts,
+    a.binaries,
+    a.datas,
+    [],
+    name='relaymd',
+    debug=False,
+    bootloader_ignore_signals=False,
+    strip=True,
+    upx=False,
+    console=True,
+    onefile=True,
+)

relaymd_cli-0.1.0/src/relaymd/cli/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"

relaymd_cli-0.1.0/src/relaymd/cli/__main__.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import sys
+
+import typer
+from relaymd.cli.commands.jobs import app as jobs_app
+from relaymd.cli.commands.submit import submit
+from relaymd.cli.commands.workers import app as workers_app
+
+cli = typer.Typer(help="RelayMD operator CLI")
+cli.command()(submit)
+cli.add_typer(jobs_app, name="jobs")
+cli.add_typer(workers_app, name="workers")
+
+
+def app() -> None:
+    cli()
+
+
+if __name__ == "__main__":
+    sys.exit(app())

relaymd_cli-0.1.0/src/relaymd/cli/commands/jobs.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+import httpx
+import typer
+from relaymd.cli.config import load_settings
+from rich.console import Console
+from rich.table import Table
+
+app = typer.Typer(help="Manage jobs.")
+console = Console()
+
+
+def _headers() -> dict[str, str]:
+    return {"X-API-Token": load_settings().api_token}
+
+
+def _orchestrator_base() -> str:
+    return load_settings().orchestrator_url.rstrip("/")
+
+
+def _status_style(status: str) -> str:
+    styles = {
+        "queued": "yellow",
+        "assigned": "cyan",
+        "running": "blue",
+        "completed": "green",
+        "failed": "red",
+        "cancelled": "dim",
+    }
+    return styles.get(status, "white")
+
+
+def _parse_iso_datetime(value: str | None) -> datetime | None:
+    if value is None:
+        return None
+    normalized = value.replace("Z", "+00:00")
+    parsed = datetime.fromisoformat(normalized)
+    return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=UTC)
+
+
+def _age_label(timestamp: str | None) -> str:
+    parsed = _parse_iso_datetime(timestamp)
+    if parsed is None:
+        return "-"
+    seconds = int((datetime.now(UTC) - parsed).total_seconds())
+    if seconds < 60:
+        return f"{seconds}s"
+    if seconds < 3600:
+        return f"{seconds // 60}m"
+    return f"{seconds // 3600}h"
+
+
+def _request(method: str, path: str, **kwargs: Any) -> httpx.Response:
+    with httpx.Client(timeout=30.0) as client:
+        response = client.request(
+            method,
+            f"{_orchestrator_base()}{path}",
+            headers=_headers(),
+            **kwargs,
+        )
+        response.raise_for_status()
+        return response
+
+
+@app.command("list")
+def list_jobs() -> None:
+    try:
+        jobs = _request("GET", "/jobs").json()
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to list jobs:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    table = Table(title="Jobs")
+    table.add_column("ID")
+    table.add_column("Title")
+    table.add_column("Status")
+    table.add_column("Assigned Worker")
+    table.add_column("Last Checkpoint Age")
+
+    for job in jobs:
+        status = str(job.get("status", "unknown"))
+        table.add_row(
+            str(job.get("id", "-"))[:8],
+            str(job.get("title", "-")),
+            f"[{_status_style(status)}]{status}[/{_status_style(status)}]",
+            str(job.get("assigned_worker_id") or "-")[:8],
+            _age_label(job.get("last_checkpoint_at")),
+        )
+
+    console.print(table)
+
+
+@app.command("status")
+def job_status(job_id: str) -> None:
+    try:
+        job = _request("GET", f"/jobs/{job_id}").json()
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to get job status:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    status = str(job.get("status", "unknown"))
+    rows = [
+        ("id", str(job.get("id", "-"))),
+        ("title", str(job.get("title", "-"))),
+        ("status", f"[{_status_style(status)}]{status}[/{_status_style(status)}]"),
+        ("input_bundle_path", str(job.get("input_bundle_path", "-"))),
+        ("latest_checkpoint_path", str(job.get("latest_checkpoint_path") or "-")),
+        ("last_checkpoint_at", str(job.get("last_checkpoint_at") or "-")),
+        ("assigned_worker_id", str(job.get("assigned_worker_id") or "-")),
+        ("created_at", str(job.get("created_at") or "-")),
+        ("updated_at", str(job.get("updated_at") or "-")),
+    ]
+
+    table = Table(title=f"Job {job_id}", show_header=False)
+    table.add_column("Field", style="cyan")
+    table.add_column("Value")
+    for key, value in rows:
+        table.add_row(key, value)
+    console.print(table)
+
+
+@app.command("cancel")
+def cancel_job(
+    job_id: str,
+    force: bool = typer.Option(False, "--force", help="Cancel running job."),
+) -> None:
+    try:
+        path = f"/jobs/{job_id}"
+        if force:
+            path = f"{path}?force=true"
+        _request("DELETE", path)
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to cancel job:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    console.print(f"[green]Cancelled job[/green] {job_id}")
+
+
+@app.command("requeue")
+def requeue_job(job_id: str) -> None:
+    try:
+        _request("POST", f"/jobs/{job_id}/requeue")
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to requeue job:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    console.print(f"[green]Requeued job[/green] {job_id}")

relaymd_cli-0.1.0/src/relaymd/cli/commands/submit.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import json
+import tarfile
+import tempfile
+import uuid
+from pathlib import Path
+from typing import Annotated, Any
+
+import httpx
+import typer
+from relaymd.cli.config import load_settings
+from relaymd.storage.client import StorageClient
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+
+console = Console()
+
+
+def ensure_worker_config(input_dir: Path, command: str | None, checkpoint_glob: str | None) -> None:
+    worker_json = input_dir / "relaymd-worker.json"
+    worker_toml = input_dir / "relaymd-worker.toml"
+
+    if command is not None:
+        worker_payload: dict[str, Any] = {
+            "command": command,
+            "checkpoint_glob_pattern": checkpoint_glob,
+        }
+        worker_json.write_text(f"{json.dumps(worker_payload, indent=2)}\n", encoding="utf-8")
+        return
+
+    if not worker_json.exists() and not worker_toml.exists():
+        console.print(
+            "[red]Missing worker configuration:[/red] provide --command or add "
+            "relaymd-worker.json / relaymd-worker.toml in the input directory."
+        )
+        raise typer.Exit(code=1)
+
+
+def create_bundle_archive(input_dir: Path, archive_path: Path) -> None:
+    with tarfile.open(archive_path, "w:gz") as tar:
+        for path in sorted(input_dir.rglob("*")):
+            if not path.is_file():
+                continue
+            arcname = path.relative_to(input_dir)
+            tar.add(path, arcname=str(arcname))
+
+
+def upload_bundle(local_archive: Path, b2_key: str) -> None:
+    settings = load_settings()
+    storage = StorageClient(
+        b2_endpoint_url=settings.b2_endpoint_url,
+        b2_bucket_name=settings.b2_bucket_name,
+        b2_access_key_id=settings.b2_access_key_id,
+        b2_secret_access_key=settings.b2_secret_access_key,
+        cf_worker_url="",
+        cf_bearer_token="",
+    )
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+        transient=True,
+    ) as progress:
+        task_id = progress.add_task("Uploading bundle to B2...", total=None)
+        storage.upload_file(local_archive, b2_key)
+        progress.update(task_id, description="Upload complete")
+
+
+def register_job(title: str, b2_key: str) -> str:
+    settings = load_settings()
+    headers = {"X-API-Token": settings.api_token}
+    payload = {"title": title, "input_bundle_path": b2_key}
+    with httpx.Client(timeout=30.0) as client:
+        response = client.post(
+            f"{settings.orchestrator_url.rstrip('/')}/jobs",
+            json=payload,
+            headers=headers,
+        )
+        response.raise_for_status()
+        return str(response.json()["id"])
+
+
+def submit(
+    input_dir: Annotated[Path, typer.Argument(help="Input directory to pack and submit.")],
+    title: Annotated[str, typer.Option("--title", help="Human-readable job title.")],
+    command: Annotated[
+        str | None,
+        typer.Option("--command", help="Command to write to relaymd-worker.json."),
+    ] = None,
+    checkpoint_glob: Annotated[
+        str | None,
+        typer.Option(
+            "--checkpoint-glob",
+            help="Checkpoint glob pattern to write alongside --command.",
+        ),
+    ] = None,
+) -> None:
+    if not input_dir.exists() or not input_dir.is_dir():
+        console.print(
+            "[red]Input directory does not exist or is not a directory:[/red] "
+            f"{input_dir}"
+        )
+        raise typer.Exit(code=1)
+
+    ensure_worker_config(input_dir, command, checkpoint_glob)
+
+    job_id = str(uuid.uuid4())
+    b2_key = f"jobs/{job_id}/input/bundle.tar.gz"
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            archive_path = Path(tmpdir) / "bundle.tar.gz"
+            create_bundle_archive(input_dir, archive_path)
+            upload_bundle(archive_path, b2_key)
+
+        created_job_id = register_job(title, b2_key)
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to submit job:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    console.print(
+        Panel.fit(
+            f"[bold green]Job submitted[/bold green]\n[bold]{created_job_id}[/bold]",
+            title="RelayMD",
+        )
+    )

relaymd_cli-0.1.0/src/relaymd/cli/commands/workers.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+import httpx
+import typer
+from relaymd.cli.config import load_settings
+from rich.console import Console
+from rich.table import Table
+
+app = typer.Typer(help="List workers.")
+console = Console()
+
+
+def _parse_iso_datetime(value: str) -> datetime:
+    normalized = value.replace("Z", "+00:00")
+    parsed = datetime.fromisoformat(normalized)
+    return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=UTC)
+
+
+@app.command("list")
+def list_workers() -> None:
+    settings = load_settings()
+    headers = {"X-API-Token": settings.api_token}
+
+    try:
+        with httpx.Client(timeout=30.0) as client:
+            response = client.get(
+                f"{settings.orchestrator_url.rstrip('/')}/workers",
+                headers=headers,
+            )
+            response.raise_for_status()
+            workers = response.json()
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Failed to list workers:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+    table = Table(title="Workers")
+    table.add_column("ID")
+    table.add_column("Platform")
+    table.add_column("GPU Model")
+    table.add_column("GPU Count")
+    table.add_column("VRAM (GB)")
+    table.add_column("Last Heartbeat")
+    table.add_column("Status")
+
+    now = datetime.now(UTC)
+    for worker in workers:
+        last_heartbeat_raw = str(worker.get("last_heartbeat", ""))
+        last_heartbeat = _parse_iso_datetime(last_heartbeat_raw)
+        age_seconds = (now - last_heartbeat).total_seconds()
+        status = "stale" if age_seconds > 120 else "healthy"
+        status_color = "red" if status == "stale" else "green"
+
+        table.add_row(
+            str(worker.get("id", "-"))[:8],
+            str(worker.get("platform", "-")),
+            str(worker.get("gpu_model", "-")),
+            str(worker.get("gpu_count", "-")),
+            str(worker.get("vram_gb", "-")),
+            last_heartbeat_raw,
+            f"[{status_color}]{status}[/{status_color}]",
+        )
+
+    console.print(table)

relaymd_cli-0.1.0/src/relaymd/cli/config.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from pydantic import AliasChoices, Field
+from pydantic_settings import (
+    BaseSettings,
+    EnvSettingsSource,
+    SettingsConfigDict,
+    YamlConfigSettingsSource,
+)
+from pydantic_settings.sources import DefaultSettingsSource, PydanticBaseSettingsSource
+
+RELAYMD_CONFIG_ENV_VAR = "RELAYMD_CONFIG"
+DEFAULT_RELAYMD_CONFIG_PATH = "~/.config/relaymd/config.yaml"
+
+
+class CliSettings(BaseSettings):
+    orchestrator_url: str = "http://localhost:8000"
+    api_token: str = Field(
+        default="change-me",
+        validation_alias=AliasChoices("api_token", "RELAYMD_API_TOKEN", "API_TOKEN"),
+    )
+    b2_endpoint_url: str = ""
+    b2_bucket_name: str = ""
+    b2_access_key_id: str = ""
+    b2_secret_access_key: str = ""
+
+    model_config = SettingsConfigDict(env_prefix="", extra="ignore")
+
+    @classmethod
+    def config_path(cls) -> Path:
+        return Path(os.getenv(RELAYMD_CONFIG_ENV_VAR, DEFAULT_RELAYMD_CONFIG_PATH)).expanduser()
+
+    @classmethod
+    def _drop_yaml_keys_with_env_overrides(
+        cls,
+        yaml_source: YamlConfigSettingsSource,
+    ) -> None:
+        env_override_map = {
+            "orchestrator_url": ("ORCHESTRATOR_URL",),
+            "api_token": ("RELAYMD_API_TOKEN", "API_TOKEN"),
+            "b2_endpoint_url": ("B2_ENDPOINT_URL",),
+            "b2_bucket_name": ("B2_BUCKET_NAME",),
+            "b2_access_key_id": ("B2_ACCESS_KEY_ID",),
+            "b2_secret_access_key": ("B2_SECRET_ACCESS_KEY",),
+        }
+        for field_name, env_keys in env_override_map.items():
+            if any(os.getenv(env_key) is not None for env_key in env_keys):
+                yaml_source.init_kwargs.pop(field_name, None)
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        yaml_source = YamlConfigSettingsSource(settings_cls, yaml_file=cls.config_path())
+        cls._drop_yaml_keys_with_env_overrides(yaml_source)
+        return (
+            init_settings,
+            yaml_source,
+            EnvSettingsSource(settings_cls),
+            DefaultSettingsSource(settings_cls),
+        )
+
+
+def load_settings() -> CliSettings:
+    return CliSettings()

relaymd_cli-0.1.0/tests/test_submit.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import json
+import tarfile
+import uuid
+from pathlib import Path
+from unittest.mock import Mock
+
+import pytest
+import typer
+from relaymd.cli.commands import submit as submit_cmd
+
+
+class _FakeSettings:
+    orchestrator_url = "http://localhost:8000"
+    api_token = "token"
+    b2_endpoint_url = "https://b2.example"
+    b2_bucket_name = "bucket"
+    b2_access_key_id = "key"
+    b2_secret_access_key = "secret"
+
+
+def test_create_bundle_archive_uses_flat_archive_root(tmp_path: Path) -> None:
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    (input_dir / "relaymd-worker.json").write_text('{"command": "run"}\n', encoding="utf-8")
+    (input_dir / "input.txt").write_text("data", encoding="utf-8")
+    nested = input_dir / "nested"
+    nested.mkdir()
+    (nested / "checkpoint.dat").write_text("cpt", encoding="utf-8")
+
+    bundle = tmp_path / "bundle.tar.gz"
+    submit_cmd.create_bundle_archive(input_dir, bundle)
+
+    with tarfile.open(bundle, "r:gz") as tar:
+        names = tar.getnames()
+
+    assert "relaymd-worker.json" in names
+    assert "input.txt" in names
+    assert "nested/checkpoint.dat" in names
+    assert all(not name.startswith("input/") for name in names)
+
+
+def test_submit_writes_worker_json_when_command_flag_provided(monkeypatch, tmp_path: Path) -> None:
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    (input_dir / "hello.txt").write_text("hello", encoding="utf-8")
+
+    uploaded = {}
+
+    class FakeStorageClient:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        def upload_file(self, local_path: Path, b2_key: str) -> None:
+            uploaded["path"] = Path(local_path)
+            uploaded["key"] = b2_key
+            with tarfile.open(local_path, "r:gz") as tar:
+                uploaded["tar_names"] = tar.getnames()
+                worker_json = tar.extractfile("relaymd-worker.json")
+                assert worker_json is not None
+                uploaded["worker_config"] = json.loads(worker_json.read().decode("utf-8"))
+
+    fake_http_client = Mock()
+    fake_response = Mock()
+    fake_response.json.return_value = {"id": str(uuid.uuid4())}
+    fake_response.raise_for_status.return_value = None
+    fake_http_client.post.return_value = fake_response
+    fake_http_cm = Mock()
+    fake_http_cm.__enter__ = Mock(return_value=fake_http_client)
+    fake_http_cm.__exit__ = Mock(return_value=False)
+
+    monkeypatch.setattr(submit_cmd, "StorageClient", FakeStorageClient)
+    monkeypatch.setattr(submit_cmd.httpx, "Client", Mock(return_value=fake_http_cm))
+    monkeypatch.setattr(submit_cmd, "load_settings", lambda: _FakeSettings())
+
+    submit_cmd.submit(
+        input_dir=input_dir,
+        title="test-job",
+        command="python run.py",
+        checkpoint_glob="*.cpt",
+    )
+
+    worker_json = input_dir / "relaymd-worker.json"
+    assert worker_json.exists()
+    worker_payload = json.loads(worker_json.read_text(encoding="utf-8"))
+    assert worker_payload == {
+        "command": "python run.py",
+        "checkpoint_glob_pattern": "*.cpt",
+    }
+    assert "relaymd-worker.json" in uploaded["tar_names"]
+    assert uploaded["worker_config"] == worker_payload
+    assert uploaded["key"].startswith("jobs/")
+    assert uploaded["key"].endswith("/input/bundle.tar.gz")
+
+
+def test_submit_aborts_when_worker_config_missing_and_no_command(tmp_path: Path) -> None:
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    (input_dir / "hello.txt").write_text("hello", encoding="utf-8")
+
+    with pytest.raises(typer.Exit) as exc:
+        submit_cmd.ensure_worker_config(input_dir, command=None, checkpoint_glob=None)
+
+    assert exc.value.exit_code == 1