| """Behavioral checks for billing, isolation, output preservation and preset recovery.""" |
|
|
| import hashlib |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| from types import SimpleNamespace |
| from unittest.mock import Mock |
|
|
| import pyarrow as pa |
| import pytest |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| import studio |
| from model_catalog import MODELS, model_profile |
| from provider import Budget, BudgetExceededError |
| from worker import FALLBACK, PreservingWriter, TunedPlan |
|
|
| from lerobot.annotations.steerable_pipeline.config import PlanConfig |
| from lerobot.annotations.steerable_pipeline.modules.plan_subtasks_memory import PlanSubtasksMemoryModule |
| from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient |
|
|
|
|
| @pytest.mark.parametrize( |
| "features", [list(studio.FEATURES), ["Interjections & speech"], ["Memory"], ["Visual questions"]] |
| ) |
| @pytest.mark.parametrize("custom", [False, True]) |
| def test_real_video_pipeline_and_checkbox_filtering(tmp_path, features, custom): |
| import av |
| import numpy as np |
| from worker import annotate_episode |
|
|
| from lerobot.annotations.steerable_pipeline.reader import EpisodeRecord |
| from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging |
|
|
| camera = "observation.images.top" |
| video = tmp_path / "videos" / camera / "chunk-000" / "file-000.mp4" |
| video.parent.mkdir(parents=True) |
| with av.open(str(video), "w") as container: |
| stream = container.add_stream("mpeg4", rate=10) |
| stream.width, stream.height, stream.pix_fmt = 64, 64, "yuv420p" |
| for i in range(40): |
| frame = av.VideoFrame.from_ndarray(np.full((64, 64, 3), i * 5, dtype=np.uint8), format="rgb24") |
| for packet in stream.encode(frame): |
| container.mux(packet) |
| for packet in stream.encode(): |
| container.mux(packet) |
| info = {"fps": 10, "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"} |
| row = { |
| "length": 40, |
| f"videos/{camera}/chunk_index": 0, |
| f"videos/{camera}/file_index": 0, |
| f"videos/{camera}/from_timestamp": 0.0, |
| } |
| record = EpisodeRecord( |
| 0, |
| "move the cup to the table", |
| tuple(i / 10 for i in range(40)), |
| tuple(range(40)), |
| tmp_path / "unused.parquet", |
| 0, |
| 40, |
| ) |
|
|
| from prompt_config import TEMPLATES, validate_prompts |
|
|
| overrides = ( |
| validate_prompts( |
| {key: item["default"] + f"\nCUSTOM_{key}" for key, item in TEMPLATES.items()}, features |
| ) |
| if custom |
| else {} |
| ) |
| seen = set() |
|
|
| def respond(messages): |
| text = " ".join(p.get("text", "") for m in messages for p in m["content"] if isinstance(p, dict)) |
| for key in overrides: |
| if f"CUSTOM_{key}" in text: |
| seen.add(key) |
| if "CUSTOM_plan_render" in text: |
| return {"plan": "1. Put the cup on the table."} |
| if "meaningful action phases" in text: |
| return { |
| "subtasks": [ |
| {"text": "grasp the cup", "start": 0, "end": 2.5}, |
| {"text": "place the cup on table", "start": 2.5, "end": 3.9}, |
| ] |
| } |
| if "compressed semantic memory" in text: |
| return {"memory": "grasped the cup"} |
| if "acknowledgement the robot" in text: |
| return {"text": "On it."} |
| if "Write ONE compact interjection" in text: |
| return {"interjection": "place the cup on the table", "speech": "Okay."} |
| if "frame-grounded visual question" in text: |
| return {"question": "How many cups?", "answer": {"label": "cup", "count": 1}} |
| if "rephras" in text.lower(): |
| return {"rephrasings": ["put the cup on the table"]} |
| raise AssertionError("Unexpected pipeline prompt") |
|
|
| result = annotate_episode( |
| record, |
| row, |
| info, |
| tmp_path, |
| {"camera": camera, "features": features, "vqa_hz": 0.1, "prompts": overrides}, |
| StubVlmClient(respond), |
| tmp_path / "staging", |
| ) |
| staged = EpisodeStaging(tmp_path / "staging", 0).read_all() |
| actual_styles = {r.get("style") for rows in staged.values() for r in rows} |
| expected = {studio.FEATURES[f] for f in features if f != studio.HUMAN_VIDEOS} |
| if "Interjections & speech" in features: |
| expected.add(None) |
| assert actual_styles == expected |
| assert result["status"] == "completed" |
| assert result["rows"] > 0 |
| assert set(overrides) == seen |
| from worker import CUSTOM_PROMPTS |
|
|
| assert CUSTOM_PROMPTS.get() is None |
|
|
|
|
| def selection(): |
| return { |
| "episodes": [{"episode_index": 0, "length": 300}, {"episode_index": 1, "length": 600}], |
| "seconds": 30, |
| "info": {"fps": 30}, |
| "dataset": "source/robot", |
| "revision": "a" * 40, |
| "start": 0, |
| "count": 2, |
| } |
|
|
|
|
| def plan(): |
| return { |
| "selection": selection(), |
| "features": ["Subtasks"], |
| "camera": "observation.images.top", |
| "vqa_hz": 0.1, |
| "mode": studio.COPY, |
| "created": time.time(), |
| "nonce": "request123", |
| "costs": studio.estimate(selection(), ["Subtasks"], 0.1, 4, 0.0005), |
| } |
|
|
|
|
| def api_mock(): |
| api = Mock() |
| api.whoami.return_value = {"name": "alice"} |
| api.list_jobs.return_value = [] |
| api.repo_exists.return_value = False |
| api.dataset_info.return_value = SimpleNamespace(id="source/robot", sha="a" * 40) |
| api.space_info.return_value = SimpleNamespace(sha="b" * 40) |
| api.list_jobs_hardware.return_value = [ |
| SimpleNamespace(name="cpu-upgrade", unit_label="minute", unit_cost_usd=0.0005) |
| ] |
| api.model_info.return_value = SimpleNamespace( |
| inference_provider_mapping={studio.PROVIDER: SimpleNamespace(status="live")} |
| ) |
| api.run_job.return_value = SimpleNamespace(id="job123") |
| return api |
|
|
|
|
| def test_cost_features_and_timeout(): |
| base = studio.estimate(selection(), ["Subtasks"], 0.1, 1, 0.0005) |
| plans = studio.estimate(selection(), ["Subtasks", "Plans"], 0.1, 1, 0.0005) |
| memory = studio.estimate(selection(), ["Subtasks", "Memory"], 0.1, 1, 0.0005) |
| vqa = studio.estimate(selection(), ["Visual questions"], 1.0, 4, 0.0005) |
| assert base["inference"] == plans["inference"] < memory["inference"] |
| assert base["cpu_max"] < vqa["cpu_max"] |
| assert base["low"] < base["high"] |
| assert vqa["calls"] == 32 |
|
|
|
|
| @pytest.mark.parametrize("value", ["../secret", "/tmp/file", "videos/../../secret", "a\\b"]) |
| def test_reject_metadata_path_traversal(value): |
| with pytest.raises(ValueError): |
| studio.relative_path(value) |
|
|
|
|
| def test_budget_unknown_failure_is_never_refunded(): |
| budget = Budget(0.09, 0.15, 0.5) |
| first = budget.reserve(4096) |
| budget.reconcile(first, None) |
| second = budget.reserve(4096) |
| budget.reconcile(second, {"prompt_tokens": 1000, "completion_tokens": 100}) |
| assert budget.reserved == pytest.approx(first + 0.0002) |
| budget.reserve(4096) |
| with pytest.raises(BudgetExceededError): |
| budget.reserve(4096) |
|
|
|
|
| @pytest.mark.parametrize("mapping_format", ["dict", "list"]) |
| def test_launch_identity_secret_and_immutable_manifest(mapping_format): |
| api = api_mock() |
| if mapping_format == "list": |
| api.model_info.return_value.inference_provider_mapping = [ |
| SimpleNamespace(provider=studio.PROVIDER, status="live") |
| ] |
| configured = plan() |
| configured["prompts"] = {"plan_subtasks": "custom instructions preserved verbatim"} |
| result = studio.launch( |
| configured, |
| "alice/annotations", |
| 5, |
| True, |
| "sensitive-token", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| ) |
| args = api.run_job.call_args.kwargs |
| assert args["namespace"] == "alice" |
| assert args["labels"]["annotation_output"] == "annotations" |
| assert all("/" not in value for value in args["labels"].values()) |
| assert args["secrets"] == {"HF_TOKEN": "sensitive-token"} |
| assert "sensitive-token" not in json.dumps(args["env"]) + json.dumps(args["command"]) |
| manifest = json.loads(args["env"]["ANNOTATION_MANIFEST"]) |
| assert manifest["episode_ids_sha256"] == hashlib.sha256(json.dumps([0, 1]).encode()).hexdigest() |
| assert manifest["space_revision"] == "b" * 40 |
| assert manifest["prompts"] == configured["prompts"] |
| assert manifest["model"] == "Qwen/Qwen3.8-27B:novita" |
| assert manifest["request_concurrency"] == manifest["episode_parallelism"] == 32 |
| assert manifest["costs"]["context_length"] == 1_000_000 |
| assert "prompt-templates.json" in str(args["command"]) |
| assert manifest["inference_budget"] + plan()["costs"]["cpu_max"] == 5 |
| assert result["id"] == "job123" |
|
|
|
|
| @pytest.mark.parametrize( |
| "output,consent,token", |
| [("bob/output", True, "token"), ("alice/output", False, "token"), ("alice/output", True, False)], |
| ) |
| def test_launch_rejects_unauthorized_or_unconfirmed(output, consent, token): |
| api = api_mock() |
| with pytest.raises(ValueError): |
| studio.launch( |
| plan(), output, 5, consent, token, space="lerobot/annotation-studio", api_factory=lambda **_: api |
| ) |
| api.run_job.assert_not_called() |
|
|
|
|
| def test_duplicate_click_recovers_existing_job(): |
| api = api_mock() |
| api.list_jobs.return_value = [ |
| SimpleNamespace(id="first-job", labels={"annotation_request": "request123"}) |
| ] |
| assert ( |
| studio.launch(plan(), "alice/output", 5, True, "token", api_factory=lambda **_: api)["id"] |
| == "first-job" |
| ) |
| api.run_job.assert_not_called() |
|
|
|
|
| def test_source_changes_block_launch(): |
| api = api_mock() |
| api.dataset_info.return_value.sha = "c" * 40 |
| with pytest.raises(ValueError, match="source changed"): |
| studio.launch(plan(), "alice/output", 5, True, "token", api_factory=lambda **_: api) |
| api.run_job.assert_not_called() |
|
|
|
|
| def test_compact_quote_preserves_pinned_episode_selection(): |
| api = api_mock() |
| request = plan() |
| digest = hashlib.sha256(json.dumps([0, 1]).encode()).hexdigest() |
| request["selection"].pop("episodes") |
| request["selection"]["episode_ids_sha256"] = digest |
| studio.launch( |
| request, |
| "alice/output", |
| 5, |
| True, |
| "token", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| ) |
| manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"]) |
| assert manifest["episode_ids_sha256"] == digest |
|
|
|
|
| def test_preserve_unselected_episodes_and_styles(): |
| def row(text, style): |
| return { |
| "role": "assistant", |
| "content": text, |
| "style": style, |
| "timestamp": 0.0, |
| "camera": None, |
| "tool_calls": None, |
| } |
|
|
| old = [[row("old subtask", "subtask"), row("keep memory", "memory")], [row("untouched", "subtask")]] |
| table = pa.table({"episode_index": [0, 1], "timestamp": [0.0, 0.0], "language_persistent": old}) |
| writer = PreservingWriter([0], {"subtask"}) |
| result = writer._materialize_table(table, [[row("new subtask", "subtask")], []], [[], []], drop_old=False) |
| rows = result.column("language_persistent").to_pylist() |
| assert {r["content"] for r in rows[0]} == {"new subtask", "keep memory"} |
| assert rows[1] == old[1] |
|
|
|
|
| def test_recovery_uses_original_prompt_and_does_not_swallow_budget(monkeypatch): |
| calls = [] |
|
|
| def generate(self, record, task=None): |
| calls.append(FALLBACK.get()) |
| if not FALLBACK.get(): |
| raise ValueError("Invalid model response") |
| return [{"text": "grasp cup", "start": 0, "end": 1}] |
|
|
| monkeypatch.setattr(PlanSubtasksMemoryModule, "_generate_subtasks", generate) |
| model = TunedPlan(vlm=StubVlmClient(lambda _: {}), config=PlanConfig()) |
| assert model._generate_subtasks(None)[0]["text"] == "grasp cup" |
| assert calls == [False, True] and not FALLBACK.get() |
| monkeypatch.setattr( |
| PlanSubtasksMemoryModule, "_generate_subtasks", Mock(side_effect=BudgetExceededError()) |
| ) |
| with pytest.raises(BudgetExceededError): |
| model._generate_subtasks(None) |
| assert not model.used_fallback |
|
|
|
|
| def test_prompt_defaults_and_validation(): |
| from prompt_config import TEMPLATES, validate_prompts |
| from worker import ACTION_PROMPT, ORIGINAL_LOAD |
|
|
| for key, item in TEMPLATES.items(): |
| default = item["default"] |
| assert validate_prompts({key: default}) == {} |
| if key == "plan_subtasks": |
| assert default == ACTION_PROMPT |
| elif key != "plan_render": |
| assert default == ORIGINAL_LOAD(key) |
| default = TEMPLATES["plan_subtasks"]["default"] |
| for bad in [ |
| "", |
| default.replace("{episode_task}", ""), |
| default + "{episode_task.__class__}", |
| default + "{unknown}", |
| default + "{episode_task:999999999}", |
| "a" * 20001, |
| ]: |
| with pytest.raises(ValueError): |
| validate_prompts({"plan_subtasks": bad}) |
| with pytest.raises(ValueError): |
| validate_prompts({"unknown": "text"}) |
| edits = {"plan_subtasks": default + "\nUse concise labels."} |
| assert validate_prompts(edits, ["Memory"]) == edits |
| assert validate_prompts(edits, ["Visual questions"]) == {} |
|
|
|
|
| def test_custom_plan_adds_estimated_calls_and_reset_restores_baseline(): |
| from prompt_config import TEMPLATES |
|
|
| features = ["Plans"] |
| base = studio.estimate(selection(), features, 0.1, 1, 0.0005) |
| template = TEMPLATES["plan_render"]["default"] |
| custom = studio.estimate( |
| selection(), features, 0.1, 1, 0.0005, {"plan_render": template + "\nBe concise."} |
| ) |
| reset = studio.estimate(selection(), features, 0.1, 1, 0.0005, {"plan_render": template}) |
| assert custom["calls"] > base["calls"] |
| assert custom["inference"] > base["inference"] |
| assert custom["custom_prompts"] == 1 |
| assert reset == base |
|
|
|
|
| def test_custom_prompt_context_resets_on_failure_and_is_thread_isolated(monkeypatch): |
| from concurrent.futures import ThreadPoolExecutor |
| from threading import Barrier |
|
|
| import worker |
| from prompt_config import TEMPLATES |
|
|
| barrier = Barrier(2) |
|
|
| def fail(*args): |
| barrier.wait(timeout=5) |
| text = worker.load_prompt("plan_subtasks") |
| fallback = FALLBACK.set(True) |
| try: |
| assert worker.load_prompt("plan_subtasks") == text |
| finally: |
| FALLBACK.reset(fallback) |
| raise ValueError(text.splitlines()[-1]) |
|
|
| monkeypatch.setattr(worker, "_annotate_episode", fail) |
|
|
| def run(name): |
| custom = TEMPLATES["plan_subtasks"]["default"] + "\n" + name |
| with pytest.raises(ValueError, match=name): |
| worker.annotate_episode( |
| None, |
| None, |
| None, |
| None, |
| {"features": ["Subtasks"], "prompts": {"plan_subtasks": custom}}, |
| None, |
| None, |
| ) |
| assert worker.CUSTOM_PROMPTS.get() is None |
|
|
| with ThreadPoolExecutor(max_workers=2) as pool: |
| list(pool.map(run, ["FIRST_USER", "SECOND_USER"])) |
|
|
|
|
| def test_custom_plan_only_uses_remaining_actions_and_default_avoids_model(): |
| from prompt_config import TEMPLATES |
| from worker import CUSTOM_PROMPTS |
|
|
| record = SimpleNamespace(episode_task="move the cup") |
| spans = [{"start": 0, "text": "reach the cup"}, {"start": 1, "text": "grasp the cup"}] |
| vlm = Mock() |
| module = TunedPlan(vlm=vlm, config=PlanConfig()) |
| assert module._generate_plan(record, spans, refresh_t=1) == "1. grasp the cup" |
| vlm.generate_json.assert_not_called() |
| vlm.generate_json.return_value = [{"plan": "1. Grasp the cup."}] |
| ctx = CUSTOM_PROMPTS.set({"plan_render": TEMPLATES["plan_render"]["default"] + "\nBe brief."}) |
| try: |
| assert module._generate_plan(record, spans, refresh_t=1) == "1. Grasp the cup." |
| sent = str(vlm.generate_json.call_args) |
| assert "grasp the cup" in sent and "reach the cup" not in sent |
| finally: |
| CUSTOM_PROMPTS.reset(ctx) |
|
|
|
|
| @pytest.mark.parametrize("token", [False, "private-reader-token"]) |
| def test_camera_discovery_reads_only_pinned_info_and_uses_request_identity(tmp_path, monkeypatch, token): |
| import tempfile |
|
|
| path = tmp_path / "info.json" |
| path.write_text( |
| json.dumps( |
| { |
| "codebase_version": "v3.0", |
| "fps": 30, |
| "features": { |
| "observation.images.top": {"dtype": "video"}, |
| "observation.images.left_wrist": {"dtype": "video"}, |
| "observation.images.right_wrist": {"dtype": "video"}, |
| "observation.images.depth": {"dtype": "video"}, |
| "observation.state": {"dtype": "float32"}, |
| "observation.images.still": {"dtype": "image"}, |
| }, |
| } |
| ) |
| ) |
| api = Mock() |
| api.dataset_info.return_value = SimpleNamespace(sha="c" * 40) |
| factory = Mock(return_value=api) |
| download = Mock(return_value=str(path)) |
| monkeypatch.setattr(studio, "HfApi", factory) |
| monkeypatch.setattr(studio, "hf_hub_download", download) |
| result = studio.read_cameras("owner/dataset", token) |
| assert result["cameras"] == [ |
| "observation.images.top", |
| "observation.images.left_wrist", |
| "observation.images.right_wrist", |
| ] |
| factory.assert_called_once_with(token=token) |
| api.list_repo_tree.assert_not_called() |
| download.assert_called_once() |
| args = download.call_args |
| assert args.args == ("owner/dataset", "meta/info.json") |
| assert args.kwargs["revision"] == "c" * 40 |
| assert args.kwargs["token"] == token |
| assert args.kwargs["cache_dir"].startswith(tempfile.gettempdir()) |
| assert not Path(args.kwargs["cache_dir"]).exists() |
|
|
|
|
| def test_camera_discovery_rejects_unsupported_sources(): |
| with pytest.raises(ValueError, match="RGB video camera"): |
| studio.video_cameras({"codebase_version": "v3.0", "fps": 30, "features": {}}) |
| with pytest.raises(ValueError, match="v3"): |
| studio.video_cameras({"codebase_version": "v2.1"}) |
|
|
|
|
| def test_usage_costs_use_reported_runtime_and_keep_unknown_charges_separate(): |
| from reporting import usage_costs |
|
|
| job = SimpleNamespace(durations=SimpleNamespace(running_secs=61)) |
| result = usage_costs( |
| job, |
| { |
| "manifest": {"costs": {"cpu_rate": 0.0005}}, |
| "usage": {"measured_inference_usd": 0.2, "reserved_including_uncertain_usd": 0.24}, |
| }, |
| ) |
| assert result["cpu_usd"] == 0.001 |
| assert result["total_usd"] == pytest.approx(0.201) |
| assert result["uncertain_inference_usd"] == pytest.approx(0.04) |
| assert usage_costs(SimpleNamespace(), {})["total_usd"] is None |
|
|
|
|
| def test_visualizer_links_only_target_materialized_copy_annotations(): |
| from reporting import annotation_url |
|
|
| report = { |
| "manifest": {"mode": studio.COPY, "output": "alice/out"}, |
| "status": "partial", |
| "written_episode_ids": [12, 5], |
| } |
| assert annotation_url(report) == "https://lerobot-visualize-dataset.hf.space/alice/out/5?tab=annotations" |
| report["manifest"]["mode"] = studio.EXPORT |
| assert annotation_url(report) is None |
| report["manifest"]["mode"] = studio.COPY |
| report["written_episode_ids"] = [] |
| assert annotation_url(report) is None |
|
|
|
|
| def test_cached_usage_and_deadline_stop_before_a_paid_request(monkeypatch): |
| import provider |
|
|
| budget = Budget(1, 1, 2, cached_input_rate=0.2) |
| reserved = budget.reserve(100) |
| budget.reconcile(reserved, {"prompt_tokens": 10, "completion_tokens": 20, "cached_tokens": 6}) |
| assert budget.summary()["measured_inference_usd"] == pytest.approx(0.0000452) |
| assert budget.summary()["input_tokens"] == 10 |
| assert budget.summary()["cached_input_tokens"] == 6 |
| post = Mock() |
| monkeypatch.setattr(provider.requests, "post", post) |
| client = provider.Provider("token", studio.MODEL, budget, deadline=time.time() - 1) |
| with pytest.raises(provider.DeadlineExceededError): |
| client.generate([[{"role": "user", "content": "hello"}]], 100, 0.2) |
| post.assert_not_called() |
|
|
|
|
| def test_five_minute_run_has_job_and_worker_deadlines(): |
| api = api_mock() |
| studio.launch( |
| plan(), |
| "alice/output", |
| 30, |
| True, |
| "token", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| timeout_minutes=5, |
| ) |
| kwargs = api.run_job.call_args.kwargs |
| assert kwargs["timeout"] == "5m" |
| assert "ANNOTATION_DEADLINE_EPOCH" in str(kwargs["command"]) |
| assert "270" in str(kwargs["command"]) |
|
|
|
|
| def test_parallel_reservations_respect_hosted_context_and_small_budget(): |
| from provider import Provider |
|
|
| budget = Budget(5, 0.42, 3, context_length=1_000_000) |
| client = Provider("token", studio.MODEL, budget, concurrency=32) |
| assert client.concurrency == 11 |
| reservations = [budget.reserve(4096) for _ in range(client.concurrency)] |
| assert budget.reserved < budget.limit |
| with pytest.raises(BudgetExceededError): |
| budget.reserve(4096) |
| budget.reconcile(reservations[0], {"prompt_tokens": 100, "completion_tokens": 40}) |
| budget.reserve(4096) |
| assert ( |
| Provider( |
| "token", studio.MODEL, Budget(30, 0.42, 3, context_length=1_000_000), concurrency=32 |
| ).concurrency |
| == 32 |
| ) |
|
|
|
|
| @pytest.mark.parametrize("key", list(MODELS)) |
| def test_selected_model_pins_route_prices_and_generation_parameters(key, monkeypatch): |
| import provider |
|
|
| profile = model_profile(key) |
| configured = plan() |
| configured["costs"] = studio.estimate(selection(), ["Subtasks"], 0.1, 1, 0.0005, model_key=key) |
| api = api_mock() |
| api.model_info.return_value.inference_provider_mapping = { |
| profile["provider"]: SimpleNamespace(status="live") |
| } |
| studio.launch( |
| configured, |
| "alice/selected", |
| 30, |
| True, |
| "secret", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| ) |
| manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"]) |
| assert manifest["model"] == profile["route"] |
| assert manifest["costs"]["input_rate"] == profile["input_rate"] |
| assert manifest["request_concurrency"] == profile["concurrency"] |
| api.model_info.assert_called_once_with(profile["model_id"], expand=["inferenceProviderMapping"]) |
| post = Mock( |
| return_value=SimpleNamespace( |
| status_code=200, |
| json=lambda: { |
| "choices": [{"message": {"content": "{}"}, "finish_reason": "stop"}], |
| "usage": {"prompt_tokens": 100, "completion_tokens": 10}, |
| }, |
| ) |
| ) |
| monkeypatch.setattr(provider.requests, "post", post) |
| client = provider.Provider( |
| "secret", |
| manifest["model"], |
| Budget(10, profile["input_rate"], profile["output_rate"]), |
| generation_options=manifest["generation_options"], |
| ) |
| client.generate([[{"role": "user", "content": "JSON please"}]], 100, 0.2) |
| body = post.call_args.kwargs["json"] |
| assert body["model"] == profile["route"] |
| assert { |
| k: v for k, v in body.items() if k not in {"model", "messages", "max_tokens", "temperature"} |
| } == profile["generation_options"] |
|
|
|
|
| def test_tiered_prices_reserve_maximum_and_reconcile_actual_prompt_tier(): |
| profile = model_profile("minimax-m3") |
| budget = Budget(10, 0.3, 1.2, 0.06, 1_000_000, profile["pricing_tiers"]) |
| assert budget.reservation_cost(1000) == pytest.approx(0.6024) |
| first = budget.reserve(1000) |
| budget.reconcile(first, {"prompt_tokens": 1000, "completion_tokens": 1000, "cached_tokens": 500}) |
| assert budget.measured == pytest.approx(0.00138) |
| second = budget.reserve(1000) |
| budget.reconcile(second, {"prompt_tokens": 524288, "completion_tokens": 1000}) |
| assert budget.measured == pytest.approx(0.00138 + 524288 * 0.6 / 1e6 + 0.0024) |
|
|
|
|
| def test_unknown_model_and_stale_model_prices_cannot_launch(): |
| with pytest.raises(ValueError, match="supported vision model"): |
| studio.prepare( |
| "source/robot", 0, 1, ["Subtasks"], None, 0.1, 1, studio.COPY, model_key="invented/model" |
| ) |
| api = api_mock() |
| configured = plan() |
| configured["costs"]["input_rate"] = 0 |
| with pytest.raises(ValueError, match="pricing changed"): |
| studio.launch( |
| configured, |
| "alice/output", |
| 30, |
| True, |
| "secret", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| ) |
| api.run_job.assert_not_called() |
|
|
|
|
| @pytest.mark.parametrize("visibility", ["private", "public"]) |
| def test_launch_pins_output_visibility(visibility): |
| api = api_mock() |
| studio.launch( |
| plan(), |
| "alice/new-output", |
| 5, |
| True, |
| "token", |
| space="lerobot/annotation-studio", |
| api_factory=lambda **_: api, |
| visibility=visibility, |
| ) |
| manifest = json.loads(api.run_job.call_args.kwargs["env"]["ANNOTATION_MANIFEST"]) |
| assert manifest["visibility"] == visibility |
|
|
|
|
| @pytest.mark.parametrize("visibility", [None, "private", "public"]) |
| def test_output_visibility_applies_to_parquet_dataset(visibility): |
| import worker |
|
|
| api = Mock() |
| api.dataset_info.return_value.sha = "pinned" |
| manifest = { |
| "dataset": "alice/source", |
| "dataset_revision": "pinned", |
| "output": "alice/new", |
| "mode": studio.COPY, |
| } |
| if visibility is not None: |
| manifest["visibility"] = visibility |
| worker.create_output_repo(api, manifest) |
| api.create_repo.assert_called_once_with( |
| "alice/new", repo_type="dataset", private=visibility != "public", exist_ok=False |
| ) |
| api.duplicate_repo.assert_not_called() |
| api.update_repo_settings.assert_not_called() |
|
|
|
|
| def test_invalid_visibility_cannot_launch_or_create_output(): |
| from worker import create_output_repo |
|
|
| api = api_mock() |
| with pytest.raises(ValueError, match="visibility"): |
| studio.launch( |
| plan(), "alice/out", 5, True, "token", api_factory=lambda **_: api, visibility="unknown" |
| ) |
| with pytest.raises(ValueError, match="visibility"): |
| create_output_repo(api, {"visibility": "unknown"}) |
| api.run_job.assert_not_called() |
| api.create_repo.assert_not_called() |
| api.duplicate_repo.assert_not_called() |
|
|
|
|
| def test_export_cannot_start_new_jobs_or_create_repositories(): |
| from worker import create_output_repo |
|
|
| api = api_mock() |
| legacy = plan() |
| legacy["mode"] = studio.EXPORT |
| with pytest.raises(ValueError, match="Parquet"): |
| studio.prepare("alice/source", 0, 1, ["Subtasks"], None, 0.1, 1, studio.EXPORT) |
| with pytest.raises(ValueError, match="Parquet"): |
| studio.launch(legacy, "alice/out", 5, True, "token", api_factory=lambda **_: api) |
| with pytest.raises(ValueError, match="Parquet"): |
| create_output_repo(api, {"mode": studio.EXPORT}) |
| api.run_job.assert_not_called() |
| api.create_repo.assert_not_called() |
|
|
