98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
import json
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
STATE_FILE = ".qc-cli.json"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CliStateStore:
|
|
config_dir: str = "."
|
|
|
|
@property
|
|
def path(self) -> Path:
|
|
return Path(self.config_dir) / STATE_FILE
|
|
|
|
def read(self) -> dict[str, Any]:
|
|
if not self.path.exists():
|
|
return {}
|
|
with open(self.path) as f:
|
|
value = json.load(f)
|
|
return dict(value) if isinstance(value, dict) else {}
|
|
|
|
def update(self, **updates: Any) -> None:
|
|
state = self.read()
|
|
state.update(updates)
|
|
self._write(state)
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
return self.read().get(key, default)
|
|
|
|
def get_last_training_job(self) -> str | None:
|
|
value = self.get("last_training_job")
|
|
return str(value) if value else None
|
|
|
|
def get_last_model_artifact(self) -> str | None:
|
|
value = self.get("last_model_artifact")
|
|
return str(value) if value else None
|
|
|
|
def get_last_quantized_model_id(self) -> str | None:
|
|
value = self.get("last_quantized_model_id")
|
|
return str(value) if value else None
|
|
|
|
def get_last_compiled_model_id(self) -> str | None:
|
|
value = self.get("last_compiled_model_id")
|
|
return str(value) if value else None
|
|
|
|
def get_last_downloaded_model(self) -> str | None:
|
|
value = self.get("last_downloaded_model")
|
|
return str(value) if value else None
|
|
|
|
def set_last_training_job(self, job_name: str) -> None:
|
|
self.update(last_training_job=job_name)
|
|
|
|
def get_training_job(self, job_name: str) -> dict[str, Any]:
|
|
jobs = self._training_jobs(self.read())
|
|
value = jobs.get(job_name, {})
|
|
return dict(value) if isinstance(value, dict) else {}
|
|
|
|
def update_training_job(self, job_name: str, **updates: Any) -> None:
|
|
state = self.read()
|
|
jobs = self._training_jobs(state)
|
|
jobs[job_name] = {**jobs.get(job_name, {}), **updates}
|
|
state["training_jobs"] = jobs
|
|
self._write(state)
|
|
|
|
def set_latest_experiment_model_version(self, version: str) -> None:
|
|
self.update(latest_experiment_model_version=version)
|
|
|
|
def get_aihub_model_provenance(self, model_id: str) -> dict[str, Any]:
|
|
provenance = self._aihub_model_provenance(self.read())
|
|
value = provenance.get(model_id, {})
|
|
return dict(value) if isinstance(value, dict) else {}
|
|
|
|
def update_aihub_model_provenance(self, model_id: str, provenance: dict[str, Any]) -> None:
|
|
state = self.read()
|
|
model_provenance = self._aihub_model_provenance(state)
|
|
model_provenance[model_id] = provenance
|
|
state["aihub_model_provenance"] = model_provenance
|
|
self._write(state)
|
|
|
|
def _write(self, state: dict[str, Any]) -> None:
|
|
with open(self.path, "w") as f:
|
|
json.dump(state, f, indent=2)
|
|
|
|
def _training_jobs(self, state: dict[str, Any]) -> dict[str, Any]:
|
|
value = state.get("training_jobs", {})
|
|
return dict(value) if isinstance(value, dict) else {}
|
|
|
|
def _aihub_model_provenance(self, state: dict[str, Any]) -> dict[str, Any]:
|
|
value = state.get("aihub_model_provenance", {})
|
|
return dict(value) if isinstance(value, dict) else {}
|
|
|
|
|
|
def store(config_path: str) -> CliStateStore:
|
|
config_dir = str(Path(config_path).parent)
|
|
return CliStateStore(config_dir)
|