modelzone-sdk 0.2.0.dev1__tar.gz → 0.3.0__tar.gz

{modelzone_sdk-0.2.0.dev1 → modelzone_sdk-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modelzone-sdk
-Version: 0.2.0.dev1
+Version: 0.3.0
 Summary: Modelzone SDK – a slim model training and serving toolkit
 License-Expression: Apache-2.0
 Author: Team Enigma
@@ -14,8 +14,11 @@ Provides-Extra: azureml
 Provides-Extra: training
 Requires-Dist: azure-ai-ml ; extra == "azureml"
 Requires-Dist: azure-identity ; extra == "azureml"
+Requires-Dist: deltalake (>=1) ; extra == "azureml"
 Requires-Dist: deltalake (>=1) ; extra == "training"
 Requires-Dist: mlflow (>=3) ; extra == "azureml"
+Requires-Dist: pandas (<3) ; extra == "azureml"
+Requires-Dist: pandas (<3) ; extra == "training"
 Description-Content-Type: text/markdown
 
 # Modelzone SDK
@@ -42,22 +45,22 @@ pip install modelzone-sdk[azureml]
 
 ## Quick start
 
-Define a model by subclassing `Model` and implementing `train` and `predict`:
+Define a model by subclassing `ModelDefinition` and implementing `train` and `predict`:
 
 ```python
-from modelzone.core import Model, PredictContext, TrainedModel
+from modelzone.core import ModelDefinition, PredictContext, ModelArtifact
 from modelzone.training import TrainingContext
 
 
-class MyModel(Model):
-    def train(self, ctx: TrainingContext) -> TrainedModel:
+class MyModel(ModelDefinition):
+    def train(self, ctx: TrainingContext) -> ModelArtifact:
         ctx.print("Training started")
 
         # … your training logic …
         fitted = train_something(seed=ctx.seed)
 
         ctx.log_metric("accuracy", 0.95)
-        return TrainedModel(model=fitted, features=["feature_a", "feature_b"])
+        return ModelArtifact(model=fitted, features=["feature_a", "feature_b"])
 
     def predict(self, ctx: PredictContext):
         df = ctx.db.query("input_table", ctx.time_interval)
@@ -91,10 +94,10 @@ result = backend.run(MyModel(), seed=42)
 ### Loading a trained model for prediction
 
 ```python
-from modelzone.predict import load_model
+from modelzone.predict import load_model_artifact
 
-trained_model = load_model(".model")
-print(trained_model.features)
+model_artifact = load_model_artifact(".model")
+print(model_artifact.features)
 ```
 
 ## CLI

{modelzone_sdk-0.2.0.dev1 → modelzone_sdk-0.3.0}/README.md

@@ -22,22 +22,22 @@ pip install modelzone-sdk[azureml]
 
 ## Quick start
 
-Define a model by subclassing `Model` and implementing `train` and `predict`:
+Define a model by subclassing `ModelDefinition` and implementing `train` and `predict`:
 
 ```python
-from modelzone.core import Model, PredictContext, TrainedModel
+from modelzone.core import ModelDefinition, PredictContext, ModelArtifact
 from modelzone.training import TrainingContext
 
 
-class MyModel(Model):
-    def train(self, ctx: TrainingContext) -> TrainedModel:
+class MyModel(ModelDefinition):
+    def train(self, ctx: TrainingContext) -> ModelArtifact:
         ctx.print("Training started")
 
         # … your training logic …
         fitted = train_something(seed=ctx.seed)
 
         ctx.log_metric("accuracy", 0.95)
-        return TrainedModel(model=fitted, features=["feature_a", "feature_b"])
+        return ModelArtifact(model=fitted, features=["feature_a", "feature_b"])
 
     def predict(self, ctx: PredictContext):
         df = ctx.db.query("input_table", ctx.time_interval)
@@ -71,10 +71,10 @@ result = backend.run(MyModel(), seed=42)
 ### Loading a trained model for prediction
 
 ```python
-from modelzone.predict import load_model
+from modelzone.predict import load_model_artifact
 
-trained_model = load_model(".model")
-print(trained_model.features)
+model_artifact = load_model_artifact(".model")
+print(model_artifact.features)
 ```
 
 ## CLI

modelzone_sdk-0.3.0/modelzone/__init__.py

@@ -0,0 +1,19 @@
+"""Modelzone SDK – structured model training and serving.
+
+Subpackages
+-----------
+- :mod:`modelzone.core` – ``ModelDefinition``, ``ModelArtifact``,
+  ``TrainingResult`` and friends (zero extras, safe everywhere).
+- :mod:`modelzone.training` – ``LocalBackend``, ``TrainingContext``,
+  ``ModelzoneDatabase`` (training environment).
+- :mod:`modelzone.azureml` – ``AzureMLBackend`` (requires ``azureml``
+  extra).
+"""
+
+from modelzone import project_config  # noqa: F401
+from modelzone.core import (  # noqa: F401
+    ModelArtifact,
+    ModelDefinition,
+    PredictContext,
+    TrainingResult,
+)

modelzone_sdk-0.3.0/modelzone/azureml/__init__.py

@@ -0,0 +1,27 @@
+"""Modelzone AzureML backend – tracked experiment runs (SDK v2 + MLflow).
+
+Requires ``azure-ai-ml``, ``azure-identity``, and ``mlflow``.
+Install with::
+
+    pip install modelzone-sdk[azureml]
+
+Usage::
+
+    from modelzone.azureml import AzureMLBackend
+
+    backend = AzureMLBackend(
+        subscription_id="...",
+        resource_group="...",
+        workspace_name="...",
+        experiment_name="my_experiment",
+    )
+    result = backend.run(model, seed=42, db=db)
+"""
+
+try:
+    from modelzone.azureml.backend import AzureMLBackend  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        "AzureMLBackend requires the 'azureml' extra. "
+        "Install it with: pip install modelzone-sdk[azureml]"
+    ) from e

{modelzone_sdk-0.2.0.dev1 → modelzone_sdk-0.3.0}/modelzone/azureml/backend.py

@@ -19,11 +19,15 @@ from typing import Any
 
 import requests
 
+from modelzone.core import (
+    ModelArtifact,
+    ModelDefinition,
+    TrainingContext,
+    TrainingResult,
+)
+from modelzone.core.constants import TRAINED_MODEL_FILENAME
 from modelzone.core.protocols import Database
-
-from modelzone.core.constants import RECORD_FILENAME, TRAINED_MODEL_FILENAME
-from modelzone.core import Model, RunRecord, TrainedModel, TrainingResult
-from modelzone.training import Backend, TrainingContext
+from modelzone.training import Backend
 
 try:
     import mlflow
@@ -85,20 +89,12 @@ class AzureMLBackend(Backend):
         self.resource_group = resource_group
         self.workspace_name = workspace_name
         self.experiment_name = experiment_name
-        self._ml_client: MLClient | None = None
-
-    @property
-    def ml_client(self) -> MLClient:
-        """Lazily connect to the AzureML workspace."""
-        if self._ml_client is None:
-            credential = DefaultAzureCredential()
-            self._ml_client = MLClient(
-                credential=credential,
-                subscription_id=self.subscription_id,
-                resource_group_name=self.resource_group,
-                workspace_name=self.workspace_name,
-            )
-        return self._ml_client
+        self.ml_client = MLClient(
+            credential=DefaultAzureCredential(),
+            subscription_id=subscription_id,
+            resource_group_name=resource_group,
+            workspace_name=workspace_name,
+        )
 
     def _configure_mlflow(self) -> None:
         """Point MLflow at the AzureML workspace tracking URI.
@@ -241,7 +237,7 @@ class AzureMLBackend(Backend):
 
     def run(
         self,
-        model: Model,
+        model: ModelDefinition,
         *,
         seed: int | None = None,
         params: dict[str, Any] | None = None,
@@ -284,8 +280,7 @@ class AzureMLBackend(Backend):
         print(f"\n  AzureML run: {portal_url}")
 
         return TrainingResult(
-            trained_model=result.trained_model,
-            record=result.record,
+            model_artifact=result.model_artifact,
             run_id=run_id,
         )
 
@@ -329,16 +324,12 @@ class AzureMLBackend(Backend):
         )
 
         with open(os.path.join(local_path, TRAINED_MODEL_FILENAME), "rb") as f:
-            trained_model: TrainedModel = pickle.load(f)  # noqa: S301
-
-        with open(os.path.join(local_path, RECORD_FILENAME), "rb") as f:
-            record: RunRecord = pickle.load(f)  # noqa: S301
+            model_artifact: ModelArtifact = pickle.load(f)  # noqa: S301
 
         shutil.rmtree(download_dir)
 
         return TrainingResult(
-            trained_model=trained_model,
-            record=record,
+            model_artifact=model_artifact,
             run_id=run_id,
         )
 
@@ -349,7 +340,7 @@ class AzureMLBackend(Backend):
     def register(
         self,
         run_id: str,
-        model_name: str,
+        registered_name: str,
         *,
         description: str | None = None,
         tags: dict[str, str] | None = None,
@@ -366,7 +357,7 @@ class AzureMLBackend(Backend):
         Args:
             run_id: The MLflow run ID whose ``run_output`` artifacts
                 will back the registered model.
-            model_name: Name for the model in the AzureML Model Registry.
+            registered_name: Name for the model in the AzureML Model Registry.
             description: Optional human-readable description.
             tags: Optional key/value tags stored on the model version.
 
@@ -380,14 +371,14 @@ class AzureMLBackend(Backend):
             merged_tags.update(tags)
 
         try:
-            existing = self.ml_client.models.list(name=model_name)
+            existing = self.ml_client.models.list(name=registered_name)
             versions = [int(m.version) for m in existing if m.version.isdigit()]
         except Exception:
             versions = []
         version = str(max(versions) + 1) if versions else "1"
 
         model = AzureMLModel(
-            name=model_name,
+            name=registered_name,
             version=version,
             path=f"azureml://jobs/{run_id}/outputs/artifacts/run_output",
             type=AssetTypes.CUSTOM_MODEL,
@@ -399,67 +390,51 @@ class AzureMLBackend(Backend):
 
     def download_registered(
         self,
-        model_name: str,
-        version: str | None = None,
+        registered_name: str,
+        version: str,
     ) -> tuple[TrainingResult, str, str]:
         """Download a registered model, returning the result and the local path.
 
         The local path points to the downloaded model directory which
-        contains ``trained_model.pkl``, ``record.pkl``, and ``code/``.
+        contains ``trained_model.pkl``, ``run_info.json``, and ``code/``.
         This is useful at build time to export the model into a
         :class:`~modelzone.training.LocalBackend` via
         :meth:`~modelzone.training.LocalBackend.save`.
 
         Args:
-            model_name: Name of the model in the AzureML Model Registry.
-            version: Version to download.  When *None* the latest
-                version is used.
+            registered_name: Name of the model in the AzureML Model Registry.
+            version: Version to download.
 
         Returns:
-            A tuple of ``(TrainingResult, model_dir_path)``.
+            A tuple of ``(TrainingResult, model_dir_path, download_dir)``.
         """
-        if version is None:
-            try:
-                existing = self.ml_client.models.list(name=model_name)
-                versions = [m for m in existing if m.version.isdigit()]
-            except Exception:
-                versions = []
-            if not versions:
-                raise ValueError(f"No versions found for model '{model_name}'")
-            latest = max(versions, key=lambda m: int(m.version))
-            version = latest.version
-
         download_dir = tempfile.mkdtemp(prefix="modelzone_registered_")
         self.ml_client.models.download(
-            name=model_name,
+            name=registered_name,
             version=version,
             download_path=download_dir,
         )
 
-        local_path = os.path.join(download_dir, model_name, "run_output")
+        local_path = os.path.join(download_dir, registered_name, "run_output")
 
         with open(os.path.join(local_path, TRAINED_MODEL_FILENAME), "rb") as f:
-            trained_model: TrainedModel = pickle.load(f)  # noqa: S301
-
-        with open(os.path.join(local_path, RECORD_FILENAME), "rb") as f:
-            record: RunRecord = pickle.load(f)  # noqa: S301
+            model_artifact: ModelArtifact = pickle.load(f)  # noqa: S301
 
         run_id = ""
-        model_entity = self.ml_client.models.get(name=model_name, version=version)
+        model_entity = self.ml_client.models.get(name=registered_name, version=version)
         if model_entity.tags:
             run_id = model_entity.tags.get("run_id", "")
 
         result = TrainingResult(
-            trained_model=trained_model,
-            record=record,
+            model_artifact=model_artifact,
             run_id=run_id,
         )
         return result, local_path, download_dir
 
     def load_registered(
         self,
-        model_name: str,
-        version: str | None = None,
+        registered_name: str,
+        version: str,
     ) -> TrainingResult:
         """Download a registered model and reconstruct a TrainingResult.
 
@@ -467,13 +442,12 @@ class AzureMLBackend(Backend):
         discards the local download path.
 
         Args:
-            model_name: Name of the model in the AzureML Model Registry.
-            version: Version to download.  When *None* the latest
-                version is used.
+            registered_name: Name of the model in the AzureML Model Registry.
+            version: Version to download.
 
         Returns:
             The reconstructed :class:`TrainingResult`.
         """
-        result, _, download_dir = self.download_registered(model_name, version)
+        result, _, download_dir = self.download_registered(registered_name, version)
         shutil.rmtree(download_dir)
         return result

{modelzone_sdk-0.2.0.dev1 → modelzone_sdk-0.3.0}/modelzone/cli.py

@@ -1,57 +1,31 @@
 """Modelzone CLI – train, register, and fetch models.
 
 Provides ``modelzone train``, ``modelzone register``, and
-``modelzone fetch`` as console entry points so that project repos
-only need config files, not scripts.
-
-All commands expect to be run from the project root directory
-containing ``project.json`` and model packages with ``model.json``.
+``modelzone fetch`` as console entry points.  All commands expect
+a ``pyproject.toml`` with a ``[tool.modelzone]`` section.
 """
 
 from __future__ import annotations
 
 import argparse
 import importlib
-import json
 import os
 import sys
 
-from modelzone import MODEL_CONFIG, PROJECT_CONFIG
-from modelzone.core import Model
-
-
-def _load_project() -> dict:
-    if not os.path.isfile(PROJECT_CONFIG):
-        print(f"Error: {PROJECT_CONFIG} not found in {os.getcwd()}", file=sys.stderr)
-        sys.exit(1)
-    with open(PROJECT_CONFIG) as f:
-        return json.load(f)
-
+from modelzone import ModelDefinition, project_config
 
-def _load_model_config(model_package: str) -> dict:
-    path = os.path.join(model_package, MODEL_CONFIG)
-    if not os.path.isfile(path):
-        print(f"Error: {path} not found in {os.getcwd()}", file=sys.stderr)
-        sys.exit(1)
-    with open(path) as f:
-        return json.load(f)
 
-
-def _load_root_model_config() -> dict:
-    if not os.path.isfile(MODEL_CONFIG):
-        print(f"Error: {MODEL_CONFIG} not found in {os.getcwd()}", file=sys.stderr)
-        sys.exit(1)
-    with open(MODEL_CONFIG) as f:
-        return json.load(f)
-
-
-def _discover_model_class(model_package: str) -> type[Model]:
+def _discover_model_class(model_package: str) -> type[ModelDefinition]:
     mod = importlib.import_module(model_package)
     for v in vars(mod).values():
-        if isinstance(v, type) and issubclass(v, Model) and v is not Model:
+        if (
+            isinstance(v, type)
+            and issubclass(v, ModelDefinition)
+            and v is not ModelDefinition
+        ):
             return v
     print(
-        f"Error: no Model subclass found in {model_package}",
+        f"Error: no ModelDefinition subclass found in {model_package}",
         file=sys.stderr,
     )
     sys.exit(1)
@@ -78,7 +52,7 @@ def _build_db(config: dict, cache: bool = True):
 
 
 def train(args: argparse.Namespace) -> None:
-    config = _load_model_config(args.model_package)
+    config = project_config.load_model(args.model_package)
     db, tags = _build_db(config, cache=not args.no_cache)
 
     params = {}
@@ -88,12 +62,12 @@ def train(args: argparse.Namespace) -> None:
     if args.azureml:
         from modelzone.azureml import AzureMLBackend
 
-        project = _load_project()
-        backend = AzureMLBackend(
+        project = project_config.load_project()
+        azureml_backend = AzureMLBackend(
             **project["workspace"],
-            experiment_name=config["experiment_name"],
+            experiment_name=config["name"],
         )
-        result = backend.run(
+        result = azureml_backend.run(
             _discover_model_class(args.model_package)(),
             seed=42,
             db=db,
@@ -105,8 +79,8 @@ def train(args: argparse.Namespace) -> None:
     else:
         from modelzone.training import LocalBackend
 
-        backend = LocalBackend(root="runs")
-        result = backend.run(
+        local_backend = LocalBackend(root="runs")
+        result = local_backend.run(
             _discover_model_class(args.model_package)(),
             seed=42,
             db=db,
@@ -116,58 +90,58 @@ def train(args: argparse.Namespace) -> None:
         )
         print(f"\nLocal Run ID: {result.run_id}")
 
-    for m in result.record.metrics:
-        print(f"  {m.name} = {m.value}")
-
 
 def register(args: argparse.Namespace) -> None:
     from modelzone.azureml import AzureMLBackend
 
-    project = _load_project()
-    config = _load_model_config(args.model_package)
+    project = project_config.load_project()
+    config = project_config.load_model(args.model_package)
 
     backend = AzureMLBackend(
         **project["workspace"],
-        experiment_name=config["experiment_name"],
+        experiment_name=config["name"],
     )
     registered = backend.register(
         run_id=args.run_id,
-        model_name=config["model_name"],
+        registered_name=config["name"],
     )
     print(f"Registered: {registered.name} v{registered.version}")
 
 
-def fetch(args: argparse.Namespace) -> None:
+def _fetch_one(
+    config: dict,
+    workspace: dict,
+) -> None:
     import pickle
     import shutil
 
     from modelzone.azureml import AzureMLBackend
 
-    project = _load_project()
-    if args.model_package:
-        config = _load_model_config(args.model_package)
-    else:
-        config = _load_root_model_config()
-    model_name = config["model_name"]
+    model_name = config["name"]
 
     backend = AzureMLBackend(
-        **project["workspace"],
-        experiment_name=config["experiment_name"],
+        **workspace,
+        experiment_name=config["name"],
     )
 
-    version = args.version or config.get("model_version")
-    version_label = str(version) if version else "latest"
-    print(f"Downloading {model_name} (version={version_label})...")
-    result, model_dir, download_dir = backend.download_registered(model_name, version=version)
+    version = config.get("model_version")
+    if version is None:
+        print(
+            f"Error: model_version is required for {model_name}",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    version = str(version)
+    print(f"Downloading {model_name} (version={version})...")
+    result, model_dir, download_dir = backend.download_registered(
+        model_name, version=version
+    )
 
-    out = args.output_dir
+    out = config["path"]
     os.makedirs(out, exist_ok=True)
 
     with open(os.path.join(out, "trained_model.pkl"), "wb") as f:
-        pickle.dump(result.trained_model, f)
-
-    with open(os.path.join(out, "record.pkl"), "wb") as f:
-        pickle.dump(result.record, f)
+        pickle.dump(result.model_artifact, f)
 
     src_code = os.path.join(model_dir, "code")
     dst_code = os.path.join(out, "code")
@@ -180,6 +154,13 @@ def fetch(args: argparse.Namespace) -> None:
     print(f"Saved to {out}")
 
 
+def fetch(args: argparse.Namespace) -> None:
+    project = project_config.load_project()
+    models = project_config.load_all_models()
+    for config in models:
+        _fetch_one(config, project["workspace"])
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(prog="modelzone", description="Modelzone CLI")
     sub = parser.add_subparsers(dest="command", required=True)
@@ -204,19 +185,9 @@ def main() -> None:
     r.add_argument("model_package", help="Model package name")
     r.add_argument("run_id", help="MLflow run ID to register")
 
-    f = sub.add_parser("fetch", help="Download a registered model for local use")
-    f.add_argument(
-        "model_package",
-        nargs="?",
-        default=None,
-        help="Model package name (omit to read model.json from cwd)",
-    )
-    f.add_argument(
-        "--version",
-        default=None,
-        help="Model version (default: model_version from config, or latest)",
+    sub.add_parser(
+        "fetch", help="Download all registered models defined in pyproject.toml"
     )
-    f.add_argument("--output-dir", default=".model", help="Local output directory")
 
     args = parser.parse_args()
     if args.command == "train":

modelzone_sdk-0.3.0/modelzone/core/__init__.py

@@ -0,0 +1,8 @@
+"""Modelzone core – minimal runtime classes for training and prediction."""
+
+from modelzone.core.constants import CODE_DIR, TRAINED_MODEL_FILENAME  # noqa: F401
+from modelzone.core.model import ModelDefinition  # noqa: F401
+from modelzone.core.model_artifact import ModelArtifact  # noqa: F401
+from modelzone.core.predict_context import PredictContext  # noqa: F401
+from modelzone.core.training_context import TrainingContext  # noqa: F401
+from modelzone.core.training_result import TrainingResult  # noqa: F401

modelzone_sdk-0.3.0/modelzone/core/constants.py

@@ -0,0 +1,9 @@
+"""Modelzone core – shared filename conventions for run directories."""
+
+# Filename used for the pickled ModelArtifact inside every run directory.
+TRAINED_MODEL_FILENAME = "trained_model.pkl"
+RUN_INFO_FILENAME = "run_info.json"
+METRICS_FILENAME = "metrics.jsonl"
+LOG_FILENAME = "output.log"
+FILES_DIR = "files"
+CODE_DIR = "code"

{modelzone_sdk-0.2.0.dev1 → modelzone_sdk-0.3.0}/modelzone/core/model.py

@@ -1,24 +1,22 @@
-"""Modelzone core – base Model class that users inherit from."""
+"""Modelzone core – base ModelDefinition class that users inherit from."""
 
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from modelzone.training.context import TrainingContext
+from typing import Any
 
+from modelzone.core.model_artifact import ModelArtifact
 from modelzone.core.predict_context import PredictContext
-from modelzone.core.trained_model import TrainedModel
+from modelzone.core.training_context import TrainingContext
 
 
-class Model:
+class ModelDefinition:
     """Base class for user-defined models.
 
     Subclass this and implement :meth:`train` and :meth:`predict`.
 
     * ``train(ctx)`` runs in the **training environment** and receives a
-      :class:`~modelzone.training.TrainingContext` for logging metrics, plots
+      :class:`~modelzone.core.TrainingContext` for logging metrics, plots
       and tags.
     * ``predict(ctx)`` runs in the **operational environment** and
       receives a :class:`~modelzone.core.PredictContext` with the
@@ -27,13 +25,13 @@ class Model:
 
     Example::
 
-        from modelzone.core import Model, TrainedModel
+        from modelzone.core import ModelDefinition, ModelArtifact
 
-        class MyModel(Model):
+        class MyModel(ModelDefinition):
             def train(self, ctx):
                 fitted = fit(load_data(), seed=ctx.seed)
                 ctx.log_metric("rmse", compute_rmse(fitted))
-                return TrainedModel(model=fitted, features=["a", "b"])
+                return ModelArtifact(model=fitted)
 
             def predict(self, ctx):
                 df = ctx.db.read("SELECT * FROM input_table")
@@ -41,8 +39,8 @@ class Model:
     """
 
     @abstractmethod
-    def train(self, ctx: TrainingContext) -> TrainedModel:
-        """Run training logic and return a :class:`TrainedModel`.
+    def train(self, ctx: TrainingContext) -> ModelArtifact:
+        """Run training logic and return a :class:`ModelArtifact`.
 
         Only called in the training environment.  Use *ctx* to log
         metrics, plots and tags; access ``ctx.seed`` and ``ctx.params``.
@@ -51,7 +49,7 @@ class Model:
             ctx: The run context for this training run.
 
         Returns:
-            A :class:`TrainedModel` with the fitted artefact(s).
+            A :class:`ModelArtifact` with the fitted artefact(s).
         """
 
     @abstractmethod