mikon 0.0.1__py3-none-any.whl

mikon/__init__.py

@@ -0,0 +1,8 @@
+from mikon.sdk.config import Config
+from mikon.sdk.context import RunContext
+from mikon.sdk import datasets
+from mikon.sdk.datasets import DatasetContext, dataset
+from mikon.sdk.job import job
+from mikon.sdk.module import ModuleFactory, ModuleRef, module
+
+__all__ = ["Config", "DatasetContext", "ModuleFactory", "ModuleRef", "RunContext", "dataset", "datasets", "job", "module"]

mikon/_runner.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import argparse
+import json
+import signal
+import sys
+import threading
+import time
+import traceback
+from pathlib import Path
+from typing import TextIO
+
+from mikon.sdk.datasets import DatasetContext, get_dataset_registry
+from mikon.sdk.context import RunContext
+from mikon.sdk.job import get_registry
+from mikon.sdk.module import instantiate_config_modules, validate_module_nest_depth
+from mikon.server.discovery import import_project
+from mikon.server.models import RunStatus
+from mikon.server.settings import load_settings
+from mikon.server.store import Store
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--run-dir", required=True)
+    args = parser.parse_args(argv)
+    run_dir = Path(args.run_dir).resolve()
+    store = Store(run_dir.parents[1])
+    meta = store.read_json(run_dir / "meta.json")
+    run_id = meta["run_id"]
+
+    stop_event = threading.Event()
+    heartbeat = threading.Thread(
+        target=_heartbeat_loop, args=(run_dir, stop_event), name="mikon-heartbeat", daemon=True
+    )
+    heartbeat.start()
+
+    def _handle_signal(signum: int, frame: object) -> None:
+        raise KeyboardInterrupt(f"received signal {signum}")
+
+    signal.signal(signal.SIGTERM, _handle_signal)
+    signal.signal(signal.SIGINT, _handle_signal)
+
+    status = RunStatus.completed
+    exit_code = 0
+    error: str | None = None
+    event_logger = _LogEventWriter(run_dir / "logs" / "events.jsonl")
+    original_stdout = sys.stdout
+    original_stderr = sys.stderr
+    sys.stdout = _TeeLineWriter(original_stdout, event_logger, "stdout")
+    sys.stderr = _TeeLineWriter(original_stderr, event_logger, "stderr")
+    try:
+        project_root = Path(meta["project_root"]).resolve()
+        watch_paths = [Path(item).resolve() for item in meta["watch"]]
+        import_project(project_root, watch_paths)
+        kind = meta.get("kind", "job")
+        if kind == "dataset":
+            definition = get_dataset_registry().get(meta["job"])
+        else:
+            definition = get_registry().get(meta["job"])
+        if definition is None:
+            raise RuntimeError(f"{kind} not found during run: {meta['job']}")
+        config_data = store.read_json(run_dir / "config.json")
+        settings = load_settings(project_root)
+        validate_module_nest_depth(config_data, settings.max_module_nest_depth)
+        config = definition.config_type.model_validate(config_data)
+        validate_module_nest_depth(config.model_dump(mode="json"), settings.max_module_nest_depth)
+        instantiate_config_modules(config, max_depth=settings.max_module_nest_depth)
+        ctx = DatasetContext(run_dir, meta["job"]) if kind == "dataset" else RunContext(run_dir)
+        definition.func(config, ctx)
+    except KeyboardInterrupt as exc:
+        status = RunStatus.stopped
+        exit_code = 143
+        error = str(exc)
+    except BaseException:
+        status = RunStatus.failed
+        exit_code = 1
+        error = traceback.format_exc()
+    finally:
+        sys.stdout = original_stdout
+        sys.stderr = original_stderr
+        event_logger.flush()
+        stop_event.set()
+        heartbeat.join(timeout=1)
+        store.write_status(run_id, status, exit_code, error)
+    return exit_code
+
+
+def _heartbeat_loop(run_dir: Path, stop_event: threading.Event) -> None:
+    heartbeat = run_dir / "heartbeat"
+    while not stop_event.is_set():
+        heartbeat.write_text(str(time.time()), encoding="utf-8")
+        stop_event.wait(2)
+
+
+class _LogEventWriter:
+    def __init__(self, path: Path) -> None:
+        self.path = path
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._seq = 0
+
+    def write_line(self, stream: str, line: str) -> None:
+        with self._lock:
+            record = {"seq": self._seq, "t": time.time(), "stream": stream, "line": line}
+            self._seq += 1
+            with self.path.open("a", encoding="utf-8") as fp:
+                fp.write(json.dumps(record, separators=(",", ":"), allow_nan=False) + "\n")
+                fp.flush()
+
+    def flush(self) -> None:
+        return
+
+
+class _TeeLineWriter:
+    def __init__(self, target: TextIO, event_logger: _LogEventWriter, stream: str) -> None:
+        self.target = target
+        self.event_logger = event_logger
+        self.stream = stream
+        self._buffer = ""
+
+    def write(self, text: str) -> int:
+        written = self.target.write(text)
+        self.target.flush()
+        self._buffer += text
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            self.event_logger.write_line(self.stream, line)
+        return written
+
+    def flush(self) -> None:
+        self.target.flush()
+        if self._buffer:
+            self.event_logger.write_line(self.stream, self._buffer)
+            self._buffer = ""
+
+    def isatty(self) -> bool:
+        return self.target.isatty()
+
+    @property
+    def encoding(self) -> str | None:
+        return self.target.encoding
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mikon/cli.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated, Any
+
+import httpx
+import typer
+import uvicorn
+
+from mikon.server.app import create_app
+from mikon.server.resources import ResourceMonitor
+from mikon.server.settings import load_settings
+
+
+app = typer.Typer(no_args_is_help=True)
+dataset_app = typer.Typer(no_args_is_help=True)
+app.add_typer(dataset_app, name="dataset")
+
+_TEMPLATES_DIR = Path(__file__).parent / "templates"
+
+
+@app.command()
+def init(
+    force: Annotated[bool, typer.Option("--force", help="Overwrite existing files.")] = False,
+) -> None:
+    root = Path.cwd()
+    _write_template(root / "mikon.toml", MIKON_TOML, force)
+    _write_template(root / "src" / "example.py", EXAMPLE_JOB, force)
+    for fname in ("USAGE.md", "USAGE-ja.md"):
+        src = _TEMPLATES_DIR / "docs" / fname
+        if src.exists():
+            _write_template(root / "docs" / fname, src.read_text(encoding="utf-8"), force)
+    typer.echo("Initialized mikon project.")
+
+
+@app.command()
+def serve(
+    host: Annotated[str, typer.Option()] = "127.0.0.1",
+    port: Annotated[int, typer.Option()] = 8000,
+    token: Annotated[str | None, typer.Option()] = None,
+) -> None:
+    if host not in {"127.0.0.1", "localhost"} and not token:
+        raise typer.BadParameter("--token is required when binding outside localhost")
+    uvicorn.run(create_app(token=token), host=host, port=port)
+
+
+@app.command("run")
+def run_job(
+    job: str,
+    gpu: Annotated[str, typer.Option("--gpu", help="Comma-separated unified GPU ids.")],
+    config: Annotated[Path | None, typer.Option("--config", exists=True, dir_okay=False)] = None,
+    set_values: Annotated[list[str] | None, typer.Option("--set", help="Override config value as key=value. Dotted keys are supported.")] = None,
+    force: Annotated[bool, typer.Option("--force")] = False,
+    server: Annotated[str, typer.Option("--server")] = "http://127.0.0.1:8000",
+) -> None:
+    config_data = _load_config(config, set_values)
+    response = httpx.post(
+        f"{server.rstrip('/')}/api/runs",
+        json={"job": job, "config": config_data, "gpus": _split_csv(gpu), "force": force},
+        timeout=30,
+    )
+    _raise_for_problem(response)
+    typer.echo(response.text)
+
+
+@app.command()
+def stop(
+    run_id: str,
+    server: Annotated[str, typer.Option("--server")] = "http://127.0.0.1:8000",
+) -> None:
+    response = httpx.post(f"{server.rstrip('/')}/api/runs/{run_id}/stop", timeout=30)
+    _raise_for_problem(response)
+    typer.echo(response.text)
+
+
+@app.command()
+def doctor() -> None:
+    settings = load_settings()
+    diagnostics = ResourceMonitor(settings).diagnostics()
+    typer.echo(diagnostics.model_dump_json(indent=2))
+
+
+@dataset_app.command("register")
+def dataset_register(
+    name: str,
+    path: Path,
+    description: Annotated[str | None, typer.Option("--description")] = None,
+    server: Annotated[str, typer.Option("--server")] = "http://127.0.0.1:8000",
+) -> None:
+    response = httpx.post(
+        f"{server.rstrip('/')}/api/datasets",
+        json={"name": name, "path": str(path), "description": description},
+        timeout=30,
+    )
+    _raise_for_problem(response)
+    typer.echo(response.text)
+
+
+@dataset_app.command("build")
+def dataset_build(
+    name: str,
+    config: Annotated[Path | None, typer.Option("--config", exists=True, dir_okay=False)] = None,
+    set_values: Annotated[list[str] | None, typer.Option("--set", help="Override config value as key=value. Dotted keys are supported.")] = None,
+    gpu: Annotated[str | None, typer.Option("--gpu", help="Comma-separated unified GPU ids.")] = None,
+    force: Annotated[bool, typer.Option("--force")] = False,
+    server: Annotated[str, typer.Option("--server")] = "http://127.0.0.1:8000",
+) -> None:
+    response = httpx.post(
+        f"{server.rstrip('/')}/api/datasets/{name}/build",
+        json={"config": _load_config(config, set_values), "gpus": _split_csv(gpu or ""), "force": force},
+        timeout=30,
+    )
+    _raise_for_problem(response)
+    typer.echo(response.text)
+
+
+def _write_template(path: Path, content: str, force: bool) -> None:
+    if path.exists() and not force:
+        typer.echo(f"Skipped existing {path}")
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content, encoding="utf-8")
+    typer.echo(f"Wrote {path}")
+
+
+def _split_csv(value: str) -> list[str]:
+    return [item.strip() for item in value.split(",") if item.strip()]
+
+
+def _load_config(config: Path | None, set_values: list[str] | None) -> dict[str, Any]:
+    config_data = json.loads(config.read_text(encoding="utf-8")) if config else {}
+    if not isinstance(config_data, dict):
+        raise typer.BadParameter("--config must contain a JSON object")
+    for item in set_values or []:
+        _apply_override(config_data, item)
+    return config_data
+
+
+def _apply_override(config: dict[str, Any], item: str) -> None:
+    if "=" not in item:
+        raise typer.BadParameter("--set values must use key=value")
+    key, raw_value = item.split("=", 1)
+    if not key:
+        raise typer.BadParameter("--set key must not be empty")
+    try:
+        value: Any = json.loads(raw_value)
+    except json.JSONDecodeError:
+        value = raw_value
+    target = config
+    parts = key.split(".")
+    for part in parts[:-1]:
+        if not part:
+            raise typer.BadParameter("--set dotted keys must not contain empty segments")
+        existing = target.get(part)
+        if existing is None:
+            existing = {}
+            target[part] = existing
+        if not isinstance(existing, dict):
+            raise typer.BadParameter(f"--set cannot assign nested key under non-object: {part}")
+        target = existing
+    if not parts[-1]:
+        raise typer.BadParameter("--set key must not end with dot")
+    target[parts[-1]] = value
+
+
+def _raise_for_problem(response: httpx.Response) -> None:
+    if response.status_code < 400:
+        return
+    try:
+        problem = response.json()
+        detail = problem.get("detail") or problem.get("title") or response.text
+    except Exception:
+        detail = response.text
+    typer.echo(f"Request failed ({response.status_code}): {detail}", err=True)
+    raise typer.Exit(1)
+
+
+MIKON_TOML = """[mikon]
+watch = ["src"]
+store = ".mikon"
+
+[gpu]
+occupancy_mem_mb = 500
+occupancy_util = 5
+
+[modules]
+max_nest_depth = 8
+
+[docs]
+root = "docs"
+"""
+
+
+EXAMPLE_JOB = '''import time
+from typing import Literal
+
+import mikon
+from mikon import Config, RunContext
+from pydantic import Field
+
+
+class ExampleConfig(Config):
+    lr: float = Field(1e-3, gt=0, le=1)
+    epochs: int = Field(5, ge=1, le=100)
+    optimizer: Literal["adam", "sgd"] = "adam"
+
+
+@mikon.job
+def example(config: ExampleConfig, ctx: RunContext) -> None:
+    for epoch in range(config.epochs):
+        loss = 1.0 / (epoch + 1)
+        print(f"epoch={epoch} loss={loss}", flush=True)
+        ctx.log_metric("loss", loss, step=epoch)
+        time.sleep(0.5)
+
+    artifact = ctx.artifacts_dir / "result.txt"
+    artifact.write_text(f"optimizer={config.optimizer}\\n", encoding="utf-8")
+    ctx.log_artifact("result.txt", artifact)
+'''

mikon/sdk/__init__.py

@@ -0,0 +1,40 @@
+from mikon.sdk.config import Config
+from mikon.sdk.context import RunContext
+from mikon.sdk import datasets
+from mikon.sdk.datasets import (
+    DatasetContext,
+    DatasetDefinition,
+    clear_dataset_registry,
+    dataset,
+    get_dataset_registry,
+)
+from mikon.sdk.job import JobDefinition, clear_registry, get_registry, job
+from mikon.sdk.module import (
+    ModuleDefinition,
+    ModuleFactory,
+    ModuleRef,
+    clear_module_registry,
+    get_module_registry,
+    module,
+)
+
+__all__ = [
+    "Config",
+    "DatasetContext",
+    "DatasetDefinition",
+    "RunContext",
+    "JobDefinition",
+    "ModuleDefinition",
+    "ModuleFactory",
+    "ModuleRef",
+    "clear_dataset_registry",
+    "clear_module_registry",
+    "clear_registry",
+    "dataset",
+    "datasets",
+    "get_dataset_registry",
+    "get_module_registry",
+    "get_registry",
+    "job",
+    "module",
+]

mikon/sdk/config.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class Config(BaseModel):
+    """Base class for job configuration schemas."""
+
+    model_config = ConfigDict(extra="forbid")

mikon/sdk/context.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import re
+import shutil
+import threading
+import time
+from pathlib import Path
+from typing import Any
+
+
+class RunContext:
+    """Runtime handle injected into mikon jobs."""
+
+    def __init__(self, run_dir: str | os.PathLike[str] | None = None) -> None:
+        raw_run_dir = run_dir or os.environ.get("MIKON_RUN_DIR")
+        if not raw_run_dir:
+            raise RuntimeError("RunContext requires run_dir or MIKON_RUN_DIR")
+        self.run_dir = Path(raw_run_dir)
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+        self.artifacts_dir = self.run_dir / "artifacts"
+        self.artifacts_dir.mkdir(parents=True, exist_ok=True)
+        self._metric_lock = threading.Lock()
+        self._artifact_lock = threading.Lock()
+        self._input_lock = threading.Lock()
+
+    def log_metric(self, name: str, value: int | float, step: int | None = None) -> None:
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("metric name must be a non-empty string")
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            raise TypeError("metric value must be int or float, but not bool")
+        if not math.isfinite(float(value)):
+            raise ValueError("metric value must be finite")
+        record: dict[str, Any] = {
+            "t": time.time(),
+            "step": step,
+            "name": name,
+            "value": float(value),
+        }
+        metrics_path = self.run_dir / "metrics.jsonl"
+        with self._metric_lock:
+            with metrics_path.open("a", encoding="utf-8") as fp:
+                fp.write(json.dumps(record, separators=(",", ":"), allow_nan=False) + "\n")
+                fp.flush()
+
+    def log_artifact(self, name: str, path: str | os.PathLike[str]) -> Path:
+        source = Path(path)
+        if not source.exists():
+            raise FileNotFoundError(source)
+
+        artifact_path = _validate_artifact_name(name)
+        destination = (self.artifacts_dir / artifact_path).resolve()
+        artifacts_root = self.artifacts_dir.resolve()
+        if not _is_relative_to(destination, artifacts_root):
+            raise ValueError("artifact name must stay within artifacts_dir")
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        if source.resolve() != destination.resolve():
+            if source.is_dir():
+                if destination.exists():
+                    shutil.rmtree(destination)
+                shutil.copytree(source, destination)
+            else:
+                shutil.copy2(source, destination)
+
+        record = {
+            "t": time.time(),
+            "name": name,
+            "path": str(destination.relative_to(self.artifacts_dir)),
+            "size": _path_size(destination),
+        }
+        with self._artifact_lock:
+            with (self.run_dir / "artifacts.jsonl").open("a", encoding="utf-8") as fp:
+                fp.write(json.dumps(record, separators=(",", ":"), allow_nan=False) + "\n")
+                fp.flush()
+        return destination
+
+    def use_dataset(self, name: str) -> Path:
+        dataset_name = _validate_name(name, "dataset")
+        meta_path = _store_root(self.run_dir) / "datasets" / dataset_name / "meta.json"
+        if not meta_path.exists():
+            raise FileNotFoundError(f"dataset not registered: {dataset_name}")
+        meta = json.loads(meta_path.read_text(encoding="utf-8"))
+        path = Path(meta["path"]).resolve()
+        if not path.exists():
+            raise FileNotFoundError(path)
+        self._log_input(
+            {
+                "type": "uses-dataset",
+                "dataset": dataset_name,
+                "path": str(path),
+            }
+        )
+        return path
+
+    def use_artifact(self, run_id: str, name: str) -> Path:
+        source_run_id = _validate_run_id(run_id)
+        artifact_path = _validate_artifact_name(name)
+        store_root = _store_root(self.run_dir)
+        artifacts_root = (store_root / "runs" / source_run_id / "artifacts").resolve()
+        path = (artifacts_root / artifact_path).resolve()
+        if not _is_relative_to(path, artifacts_root) or not path.exists():
+            raise FileNotFoundError(path)
+        self._log_input(
+            {
+                "type": "consumes-artifact",
+                "run_id": source_run_id,
+                "artifact": str(artifact_path),
+                "path": str(path),
+            }
+        )
+        return path
+
+    def _log_input(self, record: dict[str, Any]) -> None:
+        payload = {"t": time.time(), **record}
+        with self._input_lock:
+            with (self.run_dir / "inputs.jsonl").open("a", encoding="utf-8") as fp:
+                fp.write(json.dumps(payload, separators=(",", ":"), allow_nan=False) + "\n")
+                fp.flush()
+
+
+def _path_size(path: Path) -> int:
+    if path.is_file():
+        return path.stat().st_size
+    return sum(child.stat().st_size for child in path.rglob("*") if child.is_file())
+
+
+def _validate_artifact_name(name: str) -> Path:
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("artifact name must be a non-empty relative path")
+    path = Path(name)
+    if path.is_absolute() or any(part in {"", ".", ".."} for part in path.parts):
+        raise ValueError("artifact name must be a safe relative path")
+    return path
+
+
+def _validate_name(name: str, label: str) -> str:
+    if not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9_.-]+", name):
+        raise ValueError(f"{label} name must contain only letters, digits, underscore, dot, and hyphen")
+    return name
+
+
+def _validate_run_id(run_id: str) -> str:
+    if not isinstance(run_id, str) or run_id in {".", ".."} or not re.fullmatch(r"[A-Za-z0-9_.-]+", run_id):
+        raise ValueError("run_id must be a safe path segment")
+    return run_id
+
+
+def _store_root(run_dir: Path) -> Path:
+    if os.environ.get("MIKON_STORE"):
+        return Path(os.environ["MIKON_STORE"]).resolve()
+    if run_dir.parent.name == "runs":
+        return run_dir.parent.parent
+    return run_dir
+
+
+def _is_relative_to(path: Path, root: Path) -> bool:
+    try:
+        path.relative_to(root)
+        return True
+    except ValueError:
+        return False