mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 04:35:58 -10:00
Merge pull request #6244 from MervinPraison/claude/issue-6242-20261105
fix: add --no-push / ollama_push: false to keep Ollama exports local
This commit is contained in:
commit
d3ecf23533
6 files changed
+105
-2
No files matched your search
@@ -717,6 +717,9 @@ def train_export(
|
||||
mtp_draft: bool = typer.Option(
|
||||
False, "--mtp-draft/--no-mtp-draft",
|
||||
help="Also download the stock MTP drafter (Gemma-4 only) for fast inference"),
|
||||
push: Optional[bool] = typer.Option(
|
||||
None, "--push/--no-push",
|
||||
help="ollama: push to ollama.com after `ollama create` (--no-push keeps it local)"),
|
||||
):
|
||||
"""
|
||||
Export / publish an ALREADY-trained model without re-running training.
|
||||
@@ -726,6 +729,7 @@ def train_export(
|
||||
praisonai-train export gguf --model-dir lora_model --quant q4_k_m # local .gguf
|
||||
praisonai-train export gguf --model-dir lora_model --hf me/my-model # local + push
|
||||
praisonai-train export ollama --model-dir lora_model --ollama me/my-model --quant q4_k_m
|
||||
praisonai-train export ollama --model-dir lora_model --ollama my-model --no-push # local only
|
||||
"""
|
||||
from ..output.console import get_output_controller
|
||||
|
||||
@@ -765,6 +769,9 @@ def train_export(
|
||||
cfg["hf_model_name"] = hf
|
||||
if ollama:
|
||||
cfg["ollama_model"] = ollama
|
||||
# Either explicit flag overrides the config; unset keeps config/default (push).
|
||||
if isinstance(push, bool):
|
||||
cfg["ollama_push"] = push
|
||||
|
||||
# Recover the model name (drives chat-template selection) from --base-model or
|
||||
# the trained model's config.json (_name_or_path), falling back to the dir name.
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
ollama_save: "true"
|
||||
ollama_push: "true" # set "false" to keep the Ollama model local (no upload to ollama.com)
|
||||
huggingface_save: "true"
|
||||
train: "true"
|
||||
|
||||
|
||||
@@ -220,8 +220,9 @@ def create_and_push_ollama_model(
|
||||
model_parameters: str,
|
||||
modelfile_content: str,
|
||||
quantization: Optional[str] = None,
|
||||
push: bool = True,
|
||||
) -> None:
|
||||
"""Create and push an Ollama model with proper daemon management.
|
||||
"""Create (and by default push) an Ollama model with proper daemon management.
|
||||
|
||||
Args:
|
||||
ollama_model: Name of the Ollama model
|
||||
@@ -229,6 +230,8 @@ def create_and_push_ollama_model(
|
||||
modelfile_content: Content for the Modelfile
|
||||
quantization: Optional quantization (e.g. "q4_k_m") passed to
|
||||
``ollama create --quantize`` so the model isn't stored as huge f16.
|
||||
push: When False, stop after a successful ``ollama create`` and keep the
|
||||
model local-only (nothing is uploaded to ollama.com).
|
||||
|
||||
Raises:
|
||||
RuntimeError: If ollama create/push operations fail (with actionable guidance)
|
||||
@@ -259,6 +262,10 @@ def create_and_push_ollama_model(
|
||||
f"{(create.stderr or create.stdout or '').strip()}"
|
||||
)
|
||||
|
||||
if not push:
|
||||
print(f"Created local Ollama model `{tag}` (not pushed). Run it with: ollama run {tag}")
|
||||
return
|
||||
|
||||
# Push with captured output so a 401/403 becomes an actionable message rather
|
||||
# than a raw non-zero exit.
|
||||
result = subprocess.run(["ollama", "push", tag], capture_output=True, text=True)
|
||||
|
||||
@@ -545,6 +545,7 @@ class TrainModel:
|
||||
"optim", "weight_decay", "lr_scheduler_type", "seed", "output_dir",
|
||||
"assistant_only_loss", "train_on_responses_only", "save_steps",
|
||||
"train", "huggingface_save", "huggingface_save_gguf", "ollama_save",
|
||||
"ollama_push",
|
||||
"method", "beta", "max_prompt_length", "desirable_weight", "undesirable_weight",
|
||||
"reward_funcs", "num_generations", "max_completion_length",
|
||||
"embedding_learning_rate",
|
||||
@@ -1784,6 +1785,7 @@ class TrainModel:
|
||||
self.config.get('model_parameters', 'latest'),
|
||||
modelfile_content,
|
||||
quantization=self.config.get('quantization_method'),
|
||||
push=self._flag(self.config.get('ollama_push'), default=True),
|
||||
)
|
||||
|
||||
def run(self):
|
||||
|
||||
@@ -313,7 +313,8 @@ PARAMETER top_p 0.9
|
||||
create_and_push_ollama_model(
|
||||
self.config['ollama_model'],
|
||||
self.config['model_parameters'],
|
||||
modelfile_content
|
||||
modelfile_content,
|
||||
push=self._flag(self.config.get('ollama_push'), default=True),
|
||||
)
|
||||
|
||||
def run(self):
|
||||
|
||||
@@ -412,3 +412,88 @@ def test_export_command_bad_target_exits(monkeypatch, tmp_path):
|
||||
quant=None,
|
||||
base_model=None,
|
||||
)
|
||||
|
||||
|
||||
def test_ollama_no_push_skips_push(monkeypatch, tmp_path):
|
||||
"""push=False stops after `ollama create`; nothing is sent to ollama.com."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(ollama_mod, "_check_ollama_disk", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(ollama_mod, "ensure_ollama_running", lambda *_a, **_k: None)
|
||||
calls = []
|
||||
|
||||
def _fake_run(cmd, *a, **k):
|
||||
calls.append(cmd)
|
||||
return types.SimpleNamespace(returncode=0, stdout="", stderr="")
|
||||
|
||||
monkeypatch.setattr(ollama_mod.subprocess, "run", _fake_run)
|
||||
ollama_mod.create_and_push_ollama_model(
|
||||
"my-model", "latest", "FROM ./model\n", push=False)
|
||||
assert [c[1] for c in calls] == ["create"]
|
||||
assert not any("push" in c for c in calls)
|
||||
|
||||
|
||||
def _ollama_export_cfg(monkeypatch, tmp_path, push, config_push=None):
|
||||
"""Run `export ollama` with a fake trainer; return the config it received."""
|
||||
import praisonai_train.cli.commands.train as cli
|
||||
|
||||
seen = {}
|
||||
|
||||
class _FakeTrainer:
|
||||
@classmethod
|
||||
def for_export(cls, cfg):
|
||||
obj = cls()
|
||||
obj.config = cfg
|
||||
obj.model = None
|
||||
obj.hf_tokenizer = None
|
||||
return obj
|
||||
|
||||
def load_model(self):
|
||||
return ("MODEL", "TOK")
|
||||
|
||||
def create_and_push_ollama_model(self):
|
||||
seen.update(self.config)
|
||||
|
||||
monkeypatch.setattr(trainer_mod, "TrainModel", _FakeTrainer)
|
||||
config = None
|
||||
if config_push is not None:
|
||||
config = tmp_path / "cfg.yaml"
|
||||
config.write_text(f"ollama_push: {config_push}\n")
|
||||
cli.train_export(
|
||||
target="ollama", model_dir=str(tmp_path), config=str(config) if config else None,
|
||||
ollama="my-model", hf=None, quant=None, base_model="unsloth/gemma-2-2b",
|
||||
mtp_draft=False, push=push,
|
||||
)
|
||||
return seen
|
||||
|
||||
|
||||
def test_export_ollama_no_push_flag_sets_config(monkeypatch, tmp_path):
|
||||
assert _ollama_export_cfg(monkeypatch, tmp_path, push=False)["ollama_push"] is False
|
||||
|
||||
|
||||
def test_export_ollama_explicit_push_overrides_config(monkeypatch, tmp_path):
|
||||
cfg = _ollama_export_cfg(monkeypatch, tmp_path, push=True, config_push="false")
|
||||
assert cfg["ollama_push"] is True
|
||||
|
||||
|
||||
def test_export_ollama_unset_push_keeps_config(monkeypatch, tmp_path):
|
||||
cfg = _ollama_export_cfg(monkeypatch, tmp_path, push=None, config_push="false")
|
||||
assert cfg["ollama_push"] is False # YAML bool preserved, not overridden
|
||||
assert "ollama_push" not in _ollama_export_cfg(monkeypatch, tmp_path, push=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value, expected", [
|
||||
(False, False), ("false", False), (True, True), ("true", True), (None, True),
|
||||
])
|
||||
def test_llm_trainer_forwards_ollama_push(monkeypatch, value, expected):
|
||||
"""TrainModel forwards `ollama_push` (bool or string) to the helper."""
|
||||
captured = {}
|
||||
monkeypatch.setattr(
|
||||
ollama_mod, "create_and_push_ollama_model",
|
||||
lambda *a, **k: captured.update(k))
|
||||
fake = types.SimpleNamespace(
|
||||
config={"ollama_model": "m", "ollama_push": value},
|
||||
prepare_modelfile_content=lambda: "FROM ./model\n",
|
||||
_flag=trainer_mod.TrainModel._flag,
|
||||
)
|
||||
trainer_mod.TrainModel.create_and_push_ollama_model(fake)
|
||||
assert captured["push"] is expected
|
||||
Reference in new issue
Block a user