geosave-engine 0.1.0__py3-none-any.whl

geosave_engine/__about__.py

@@ -0,0 +1,2 @@
+__version__ = "0.1.0"
+__author__ = "Fatah Muria"

geosave_engine/cli/commands/__init__.py

@@ -0,0 +1 @@
+"""GeoSave CLI command modules."""

geosave_engine/cli/commands/create.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import questionary
+import typer
+
+from geosave_engine.cli.errors import AbortedByUserError
+from geosave_engine.cli.workspace import Workspace, WorkspaceSpec
+from geosave_engine.cli.workspace.templates import get_method_templates
+
+
+def create(
+    directory: Path = typer.Option(
+        Path("."),
+        "--dir",
+        "-d",
+        help="Directory to build the GeoSave workspace in.",
+    ),
+) -> None:
+    """Create one GeoSave workspace."""
+    workspace = Workspace(directory, _prompt_workspace_spec())
+    workspace.setup_workspace()
+
+
+def _prompt_workspace_spec() -> WorkspaceSpec:
+    method_templates = get_method_templates()
+    project_name = _ask_required_text("Enter the project name:", "Project name")
+    project_task = _ask_required_choice(
+        "Select the main task for the project:",
+        list(method_templates),
+        "Project task",
+    )
+    project_method = _ask_required_choice(
+        "Select the method for the project:",
+        list(method_templates[project_task]),
+        "Project method",
+    )
+    description = questionary.text("Enter a description for the project (optional):").ask()
+
+    return WorkspaceSpec(
+        project_name=project_name,
+        project_task=project_task,
+        project_method=project_method,
+        description=description.strip() if description else None,
+    )
+
+
+def _ask_required_text(question: str, field_name: str) -> str:
+    answer = questionary.text(question).ask()
+    if not answer or not answer.strip():
+        raise AbortedByUserError(f"{field_name} is required.")
+    return answer.strip()
+
+
+def _ask_required_choice(question: str, choices: list[str], field_name: str) -> str:
+    answer = questionary.select(question, choices=choices).ask()
+    if answer is None:
+        raise AbortedByUserError(f"{field_name} selection was aborted.")
+    return answer.strip()

geosave_engine/cli/commands/infra.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import importlib.resources as pkg_resources
+import shutil
+import subprocess
+from pathlib import Path
+
+import typer
+
+infra_app = typer.Typer(help="Manage GeoSave dev infrastructure (Docker Compose).")
+
+_COMPOSE_NAME = "docker-compose.yml"
+_ENV_EXAMPLE = ".env.example"
+_LOCAL_INFRA_DIR = "docker"
+
+
+def _bundled(filename: str) -> Path:
+    return Path(str(pkg_resources.files("geosave_engine") / "infra" / filename))
+
+
+def _local_compose() -> Path | None:
+    path = Path.cwd() / _LOCAL_INFRA_DIR / _COMPOSE_NAME
+    return path if path.exists() else None
+
+
+@infra_app.command()
+def init() -> None:
+    """Copy docker-compose.yml and .env to ./docker/."""
+    destination = Path.cwd() / _LOCAL_INFRA_DIR
+    destination.mkdir(exist_ok=True, parents=True)
+    compose_destination = destination / _COMPOSE_NAME
+    env_destination = destination / ".env"
+
+    if compose_destination.exists():
+        typer.echo(f"{_COMPOSE_NAME} already exists — skipped.")
+    else:
+        shutil.copy2(_bundled(_COMPOSE_NAME), compose_destination)
+        typer.echo(f"Created docker/{_COMPOSE_NAME}")
+
+    if env_destination.exists():
+        typer.echo(".env already exists — skipped.")
+    else:
+        shutil.copy2(_bundled(_ENV_EXAMPLE), env_destination)
+        typer.echo("Created docker/.env from defaults")
+
+    typer.echo("\nNext: geosave infra up")
+
+
+@infra_app.command()
+def up(
+    profile: list[str] = typer.Option([], "--profile", "-p", help="Enable a service profile."),
+    detach: bool = typer.Option(
+        True,
+        "--detach/--no-detach",
+        "-d/-D",
+        help="Run in background.",
+    ),
+) -> None:
+    """Start infrastructure containers."""
+    _run_compose(["up", *(["-d"] if detach else [])], profile)
+
+
+@infra_app.command()
+def down(profile: list[str] = typer.Option([], "--profile", "-p")) -> None:
+    """Stop infrastructure containers."""
+    _run_compose(["down"], profile)
+
+
+@infra_app.command()
+def status() -> None:
+    """Show running infrastructure containers."""
+    _run_compose(["ps"], [])
+
+
+def _run_compose(args: list[str], profiles: list[str]) -> None:
+    profile_flags = [flag for profile in profiles for flag in ("--profile", profile)]
+    compose_path = _local_compose() or _bundled(_COMPOSE_NAME)
+    command = ["docker", "compose", "-f", str(compose_path), *profile_flags, *args]
+    subprocess.run(command, check=True)

geosave_engine/cli/commands/upload.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import importlib
+import os
+import pickle
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import questionary
+import typer
+import yaml
+
+from geosave_engine.cli.errors import AbortedByUserError, WorkspaceError
+
+if TYPE_CHECKING:
+    # Only ever used as type annotations (lazy strings, thanks to
+    # `from __future__ import annotations`) — never called by name at
+    # runtime, so keep them out of the module's real import graph. Both
+    # drag in ~10s of lightning.pytorch/mlflow import time, which main.py
+    # would otherwise pay on every `geosave` invocation (create, artifact,
+    # --help, ...), not just `upload`.
+    from lightning.pytorch import LightningModule
+    from mlflow.models.model import ModelInfo
+from geosave_engine.cli.workspace import Workspace, load_run_artifact
+from geosave_engine.cli.workspace.artifact import (
+    artifact_paths,
+    resolve_artifact_name,
+    select_checkpoint,
+)
+
+
+def upload(
+    project_dir: Path | None = typer.Argument(None, help="Workspace directory."),
+    artifact_name: str | None = typer.Option(
+        None,
+        "--artifact",
+        "-a",
+        help="Artifact directory (for example, model_name/version_0).",
+    ),
+    checkpoint: str | None = typer.Option(
+        None,
+        "--checkpoint",
+        "-c",
+        help="Checkpoint filename inside the run's checkpoints/ dir. Prompt if omitted and multiple exist.",
+    ),
+    registered_model_name: str | None = typer.Option(
+        None,
+        "--name",
+        "-n",
+        help="MLflow registered model name. Defaults to the resolved run name.",
+    ),
+) -> None:
+    """Upload one trained checkpoint to the MLflow model registry.
+
+    Resolve one run artifact (discover it under artifacts/, prompt if
+    ambiguous), rebuild its LightningModule from saved config + checkpoint
+    weights, then
+    log + register it to MLflow inside a fresh mlflow.start_run(), with
+    workspace modules/ bundled as code.
+
+    Args:
+        project_dir: Workspace directory. Defaults to cwd.
+        artifact_name: Artifact directory relative to artifacts/, for
+            example "DynamicWorld/version_9". Prompt if omitted.
+        checkpoint: Checkpoint filename. Prompt if omitted and multiple
+            checkpoints exist.
+        registered_model_name: MLflow registered model name. Defaults to
+            the resolved RunArtifact.model_name.
+
+    Raises:
+        WorkspaceError: If artifact/config is missing or invalid.
+        AbortedByUserError: If prompted for a missing MLFLOW_TRACKING_URI or
+            MLFLOW_EXPERIMENT_NAME and no answer is given.
+    """
+    workspace = Workspace.load_workspace(project_dir or Path.cwd())
+
+    # modules/ has no __init__.py (namespace package) — only resolves once
+    # workspace.root is searchable. python main.py gets this for free from
+    # cwd; this installed console script doesn't, so make it explicit.
+    # No-op for premade (Path A) classes — those come from the installed
+    # geosave_engine package and resolve regardless.
+    sys.path.insert(0, str(workspace.root))
+
+    resolved_name = resolve_artifact_name(workspace, artifact_name)
+    run_dir = artifact_paths(workspace)[resolved_name]
+    artifact = load_run_artifact(run_dir)
+
+    checkpoint_path = select_checkpoint(artifact, checkpoint)
+    model = _load_model_from_checkpoint(artifact.config_path, checkpoint_path)
+
+    resolved_registered_name = registered_model_name or artifact.model_name
+    tracking_uri = _require_tracking_uri()
+    experiment_name = _require_experiment_name(resolved_registered_name)
+
+    model_info = log_model(
+        model=model,
+        name=resolved_registered_name,
+        checkpoint_path=checkpoint_path,
+        modules_dir=workspace.modules_dir,
+        tracking_uri=tracking_uri,
+        experiment_name=experiment_name,
+    )
+    # model_info.model_uri is mlflow 3.x's models:/m-<hash> model-id form —
+    # print the conventional name/version form instead, since that's what
+    # litserve configs and humans actually reference.
+    typer.echo(f"models:/{resolved_registered_name}/{model_info.registered_model_version}")
+
+
+def _import_model_class(class_path: str) -> type[LightningModule]:
+    """Dynamically import a LightningModule class from its dotted path.
+
+    Args:
+        class_path: Fully qualified class path (for example,
+            "geosave_engine.ml.tasks.SemanticSegmentationTask"), read from
+            the run's config.yaml "model.class_path".
+
+    Returns:
+        Imported class, ready for load_from_checkpoint.
+
+    Raises:
+        WorkspaceError: If the module or class can't be imported.
+    """
+    module_name, _, class_name = class_path.rpartition(".")
+    try:
+        module = importlib.import_module(module_name)
+        return getattr(module, class_name)
+    except (ImportError, AttributeError) as error:
+        raise WorkspaceError(f"Could not import model class '{class_path}': {error}") from error
+
+
+def _load_model_from_checkpoint(config_path: Path, checkpoint_path: Path) -> LightningModule:
+    """Rebuild one trained LightningModule from its run artifacts.
+
+    Uses the checkpoint's saved hyperparameters (save_hyperparameters() at
+    training time) — no need to re-pass init_args from config.yaml.
+
+    Args:
+        config_path: Run's config.yaml, used only to resolve the model
+            class_path.
+        checkpoint_path: Checkpoint file with saved weights + hparams.
+
+    Returns:
+        Instantiated model with weights loaded, ready to log.
+
+    Raises:
+        WorkspaceError: If config.yaml lacks a "model.class_path" entry, or
+            the checkpoint file is unreadable or corrupted.
+    """
+    from geosave_engine.ml.inference.protocol import Predictable
+
+    config = yaml.safe_load(config_path.read_text())
+    class_path = config.get("model", {}).get("class_path")
+    if not class_path:
+        raise WorkspaceError(f"config.yaml missing model.class_path: {config_path}")
+
+    model_cls = _import_model_class(class_path)
+    try:
+        model = model_cls.load_from_checkpoint(str(checkpoint_path), map_location="cpu")
+    except (RuntimeError, OSError, EOFError, pickle.UnpicklingError) as error:
+        raise WorkspaceError(f"Could not load checkpoint '{checkpoint_path.name}': {error}") from error
+
+    if not isinstance(model, Predictable):
+        raise WorkspaceError(
+            f"{class_path} doesn't implement Predictable (no usable predict() method) — "
+            "can't register a model that can't be served."
+        )
+    return model
+
+
+def _require_tracking_uri() -> str:
+    """Read MLFLOW_TRACKING_URI from the environment, prompting when unset.
+
+    No local-mlruns fallback — upload always targets a real registry, so an
+    empty answer aborts instead of silently defaulting.
+
+    Returns:
+        Tracking URI value.
+
+    Raises:
+        AbortedByUserError: If prompted and the user gives no answer.
+    """
+    tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
+    if tracking_uri:
+        return tracking_uri
+
+    answer = questionary.text(
+        "MLFLOW_TRACKING_URI is not set. Enter your MLflow tracking URI:"
+    ).ask()
+    if not answer or not answer.strip():
+        raise AbortedByUserError("MLflow tracking URI is required.")
+    return answer.strip()
+
+
+def _require_experiment_name(default: str) -> str:
+    """Read MLFLOW_EXPERIMENT_NAME from the environment, prompting when unset.
+
+    Mirrors GeosaveCLI's own training-time default (falls back to
+    model_name when MLFLOW_EXPERIMENT_NAME isn't set — see ml/cli/cli.py).
+
+    Args:
+        default: Prefilled answer, typically the resolved model name.
+
+    Returns:
+        Experiment name value.
+
+    Raises:
+        AbortedByUserError: If prompted and the user gives no answer.
+    """
+    experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME")
+    if experiment_name:
+        return experiment_name
+
+    answer = questionary.text(
+        "MLFLOW_EXPERIMENT_NAME is not set. Enter an experiment name:",
+        default=default,
+    ).ask()
+    if not answer or not answer.strip():
+        raise AbortedByUserError("MLflow experiment name is required.")
+    return answer.strip()
+
+
+def log_model(
+    model: LightningModule,
+    name: str,
+    checkpoint_path: Path,
+    modules_dir: Path,
+    tracking_uri: str,
+    experiment_name: str,
+) -> ModelInfo:
+    """Log one model run to MLflow, resolving tracking uri/experiment/run itself.
+
+    MLflow acts as registry using pytorch log_model. modules_dir is
+    bundled as code_paths so registry-side consumers (litserve) can import
+    workspace-local preprocessing (for example, modules.data_pipeline) after
+    pulling the model; litserve owns preprocessing + inference wiring, not
+    the logged artifact.
+
+    Args:
+        model: Instantiated model to upload.
+        name: Resolved name (for example, "DynamicWorld") — used as the
+            MLflow run name, the artifact's name inside it, and the
+            registered model name. upload only exposes one --name flag, so
+            these three never differ in practice.
+        checkpoint_path: Local checkpoint this model was rebuilt from,
+            stored as metadata for traceability.
+        modules_dir: Workspace modules/ directory bundled as code_paths —
+            code MLflow copies into the model's code/ dir and prepends to
+            sys.path on load.
+        tracking_uri: MLflow tracking URI to log against.
+        experiment_name: MLflow experiment to log under.
+
+    Returns:
+        Logged MLflow model information.
+    """
+    import mlflow
+
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.set_experiment(experiment_name)
+
+    typer.echo(f"Uploading '{name}' ({checkpoint_path.name}) to {tracking_uri} ...", err=True)
+    with mlflow.start_run(run_name=name):
+        # MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR defaults on — mlflow shows
+        # its own tqdm byte-level progress for the artifact upload itself.
+        model_info = mlflow.pytorch.log_model(
+            pytorch_model=model,
+            name=name,
+            code_paths=[str(modules_dir)],
+            registered_model_name=name,
+            metadata={"checkpoint": checkpoint_path.name},
+        )
+    typer.echo(f"Registered '{name}' version {model_info.registered_model_version}", err=True)
+
+    return model_info

geosave_engine/cli/errors.py

@@ -0,0 +1,12 @@
+class GeosaveCliError(Exception):
+    """Base error raised by GeoSave commands."""
+
+    exit_code: int = 1
+
+
+class AbortedByUserError(GeosaveCliError):
+    """The user cancelled a prompt (Ctrl-C, empty answer)."""
+
+
+class WorkspaceError(GeosaveCliError):
+    """A runtime command could not locate or load a workspace."""

geosave_engine/cli/main.py

@@ -0,0 +1,11 @@
+import typer
+
+from geosave_engine.cli.commands.create import create
+from geosave_engine.cli.commands.infra import infra_app
+from geosave_engine.cli.commands.upload import upload
+
+app = typer.Typer(help="GeoSave Engine CLI", no_args_is_help=True)
+app.command()(create)
+app.command()(upload)
+
+app.add_typer(infra_app, name="infra")

geosave_engine/cli/workspace/__init__.py

@@ -0,0 +1,12 @@
+"""Workspace loading, discovery, and scaffolding."""
+
+from .artifact import RunArtifact, discover_artifacts, load_run_artifact
+from .model import Workspace, WorkspaceSpec
+
+__all__ = [
+    "RunArtifact",
+    "Workspace",
+    "WorkspaceSpec",
+    "discover_artifacts",
+    "load_run_artifact",
+]

geosave_engine/cli/workspace/artifact.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import questionary
+
+from geosave_engine.cli.errors import AbortedByUserError, WorkspaceError
+
+if TYPE_CHECKING:
+    from .model import Workspace
+
+_CONFIG_FILE = "config.yaml"
+_CHECKPOINTS_DIR = "checkpoints"
+_CHECKPOINT_GLOB = "*.ckpt"
+
+
+@dataclass(frozen=True)
+class RunArtifact:
+    """Store paths and identity for one model run.
+
+    Args:
+        run_dir: Version directory containing artifacts for one run
+            (for example, artifacts/DynamicWorld/version_9).
+        config_path: Lightning config saved for the run.
+        checkpoint_paths: All discovered checkpoints for the run, sorted.
+    """
+
+    run_dir: Path
+    config_path: Path
+    checkpoint_paths: list[Path]
+
+    @property
+    def model_name(self) -> str:
+        """Return the run's parent directory name (for example, "DynamicWorld")."""
+        return self.run_dir.parent.name
+
+
+def discover_artifacts(artifacts_dir: Path) -> list[Path]:
+    """Find run version directories without reading their contents.
+
+    Args:
+        artifacts_dir: Workspace artifacts directory.
+
+    Returns:
+        Sorted version directories (artifacts/<model_name>/version_N) that
+        hold a config.yaml.
+    """
+    if not artifacts_dir.is_dir():
+        return []
+
+    return sorted(
+        path.resolve()
+        for path in artifacts_dir.glob("*/version_*")
+        if path.is_dir() and (path / _CONFIG_FILE).is_file()
+    )
+
+
+def load_run_artifact(run_dir: Path) -> RunArtifact:
+    """Load one run directory using the canonical artifact layout.
+
+    Args:
+        run_dir: Version directory containing one model run.
+
+    Returns:
+        Validated paths for the run.
+
+    Raises:
+        WorkspaceError: If required files are missing.
+    """
+    resolved_run_dir = run_dir.expanduser().resolve()
+    if not resolved_run_dir.is_dir():
+        raise WorkspaceError(f"Artifact run directory not found: {resolved_run_dir}")
+
+    config_path = _require_file(resolved_run_dir / _CONFIG_FILE)
+    checkpoint_paths = _discover_checkpoints(resolved_run_dir / _CHECKPOINTS_DIR)
+
+    return RunArtifact(
+        run_dir=resolved_run_dir,
+        config_path=config_path,
+        checkpoint_paths=checkpoint_paths,
+    )
+
+
+def _require_file(path: Path) -> Path:
+    if not path.is_file():
+        raise WorkspaceError(f"Required artifact file not found: {path}")
+    return path.resolve()
+
+
+def _discover_checkpoints(checkpoints_dir: Path) -> list[Path]:
+    if not checkpoints_dir.is_dir():
+        raise WorkspaceError(f"Checkpoints directory not found: {checkpoints_dir}")
+
+    checkpoint_paths = sorted(path.resolve() for path in checkpoints_dir.glob(_CHECKPOINT_GLOB))
+    if not checkpoint_paths:
+        raise WorkspaceError(f"No checkpoint files found in: {checkpoints_dir}")
+    return checkpoint_paths
+
+
+def artifact_paths(workspace: Workspace) -> dict[str, Path]:
+    """Map artifact keys to their run directories.
+
+    Args:
+        workspace: Loaded workspace to scan.
+
+    Returns:
+        Keys like "model_name/version_0" mapped to their run directory.
+    """
+    return {
+        str(path.relative_to(workspace.artifacts_dir)): path
+        for path in workspace.artifacts
+    }
+
+
+def resolve_artifact_name(workspace: Workspace, artifact_name: str | None) -> str:
+    """Resolve one artifact key, prompting when omitted.
+
+    Args:
+        workspace: Loaded workspace to scan.
+        artifact_name: Explicit artifact key, skipping the prompt.
+
+    Returns:
+        Validated artifact key, usable with artifact_paths(workspace).
+
+    Raises:
+        WorkspaceError: If no artifacts exist, or artifact_name doesn't
+            match any.
+        AbortedByUserError: If the prompt is cancelled.
+    """
+    paths = artifact_paths(workspace)
+    if artifact_name is None:
+        if not paths:
+            raise WorkspaceError(f"No artifacts found in: {workspace.artifacts_dir}")
+        artifact_name = _prompt_for_artifact(list(paths))
+
+    if artifact_name not in paths:
+        available = ", ".join(sorted(paths))
+        raise WorkspaceError(f"Artifact not found: {artifact_name}. Available: {available}")
+
+    return artifact_name
+
+
+def select_checkpoint(artifact: RunArtifact, checkpoint: str | None) -> Path:
+    """Resolve one checkpoint from a loaded RunArtifact, prompting if ambiguous.
+
+    Args:
+        artifact: Loaded run artifact with all discovered checkpoint_paths.
+        checkpoint: Checkpoint filename to use directly, skipping the
+            prompt. Prompt only when multiple checkpoints exist.
+
+    Returns:
+        Resolved checkpoint path.
+
+    Raises:
+        WorkspaceError: If checkpoint doesn't match any discovered file.
+        AbortedByUserError: If the prompt is cancelled.
+    """
+    checkpoint_paths = artifact.checkpoint_paths
+
+    if checkpoint is not None:
+        matches = [path for path in checkpoint_paths if path.name == checkpoint]
+        if not matches:
+            available = ", ".join(path.name for path in checkpoint_paths)
+            raise WorkspaceError(f"Checkpoint not found: {checkpoint}. Available: {available}")
+        return matches[0]
+
+    if len(checkpoint_paths) == 1:
+        return checkpoint_paths[0]
+
+    answer = questionary.select(
+        "Select a checkpoint:",
+        choices=[path.name for path in checkpoint_paths],
+    ).ask()
+    if answer is None:
+        raise AbortedByUserError("Checkpoint selection was aborted by the user.")
+
+    return next(path for path in checkpoint_paths if path.name == answer.strip())
+
+
+def _prompt_for_artifact(artifact_keys: list[str]) -> str:
+    answer = questionary.select("Select an artifact:", choices=artifact_keys).ask()
+    if answer is None:
+        raise AbortedByUserError("Artifact selection was aborted by the user.")
+    return answer.strip()