diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 688785aa1..26ce2406f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,7 +20,12 @@ repos: # into their working .venv. Otherwise an extra like `hardware` resolves # imports that dev-only CI cannot, basedpyright auto-prunes those baseline # entries, and the shrunk baseline makes CI flag them as new errors. - entry: env UV_PROJECT_ENVIRONMENT=.venv-typecheck uv run --locked --exact --extra dev basedpyright + # `--pythonpath` names that interpreter: basedpyright resolves imports against + # the environment it is pointed at, not the one it runs in, and left to itself + # it picks the working `.venv` — the very drift the dedicated venv rules out. + entry: >- + env UV_PROJECT_ENVIRONMENT=.venv-typecheck uv run --locked --exact --extra dev + basedpyright --pythonpath .venv-typecheck/bin/python language: system # Whole-project checker (pass_filenames: false), so run on every commit — # a pyproject/baseline-only change or a module deletion can still shift the diff --git a/positronic/cfg/analysis.py b/positronic/cfg/analysis.py index 72c1afb2a..5ddb6c67d 100644 --- a/positronic/cfg/analysis.py +++ b/positronic/cfg/analysis.py @@ -45,6 +45,11 @@ def _model_label_from_path(model_type: str, checkpoint_path: str) -> str | None: def model(ep: Episode) -> str: + # The operator's name for the endpoint that served the episode. Two deployments of one checkpoint share + # everything derived from that checkpoint, so the name outranks it. + if label := ep.get('inference.policy.label', ''): + return label + policy_type = ep.get('inference.policy.type', '') if policy_type == 'remote': diff --git a/positronic/cfg/policy.py b/positronic/cfg/policy.py index 4c612e0f4..0b10cec15 100644 --- a/positronic/cfg/policy.py +++ b/positronic/cfg/policy.py @@ -67,19 +67,23 @@ def balanced(balance: int): return BalancedSampler(balance=balance) -@cfn.config(endpoints={}, weights={}, recording_dir=None, sampler=None, group_fields=None) +@cfn.config(endpoints={}, weights={}, recording_dir=None, sampler=None, group_fields=None, headers=None) def production( endpoints: dict[str, str], weights: dict[str, float], recording_dir: str | None, sampler: Sampler | None, group_fields: list[str] | None, + headers: dict[str, str] | None, ): """Routes each episode to one of several remote endpoints, each named for CLI overrides. An endpoint is one URL, so `--policy.endpoints.groot=desktop:8000` adds or repoints one without restating the others. `weights` name the same endpoints and set their sampling odds; endpoints left - out of it weigh 1.0. + out of it weigh 1.0. `headers` reach every endpoint, since one set of credentials fronts them all. + + The endpoint's name is what identifies it — the sampling key, and the field recorded on each episode. + Two deployments of one checkpoint report the same server metadata, so only the name tells them apart. """ if not endpoints: raise ValueError('At least one endpoint must be given, e.g. --policy.endpoints.groot=desktop:8000') @@ -88,9 +92,11 @@ def production( # Every Sampler but the default uniform one picks by episode counts alone, so weights would be dropped. if weights and sampler is not None: raise ValueError(f'weights cannot be combined with {type(sampler).__name__}, which samples by count') - policies = [RemotePolicy(url, recording_dir=recording_dir) for url in endpoints.values()] + policies = [ + RemotePolicy(url, label=name, recording_dir=recording_dir, headers=headers) for name, url in endpoints.items() + ] w = [weights.get(name, 1.0) for name in endpoints] if weights else None - return SampledPolicy(*policies, weights=w, sampler=sampler, group_fields=group_fields) + return SampledPolicy(*policies, weights=w, sampler=sampler, group_fields=group_fields, key_field='label') @cfn.config() @@ -102,8 +108,33 @@ def phail_single(hostname, w_openpi=1.0, w_groot=1.0, w_act=1.0): return SampledPolicy(openpi, groot, act, weights=[w_openpi, w_groot, w_act]) +EVAL_GROUP_FIELDS = [keys.TASK, 'eval.object', 'eval.tote_placement', 'eval.external_camera'] + phail_multiple = production.override( endpoints={'smolvla': 'notebook:8000', 'act': 'notebook:8001', 'groot': 'desktop:8000', 'openpi': 'vm-openpi:8000'}, sampler=balanced, - group_fields=[keys.TASK, 'eval.object', 'eval.tote_placement', 'eval.external_camera'], + group_fields=EVAL_GROUP_FIELDS, +) + +# The blind set each Runway owner is asking the rig for right now. They ship a new batch of checkpoints every +# few days, so repoint theirs rather than adding a preset per batch, and the eval CLI stays one line. These +# deployments sit behind the workspace's proxy, so pass --policy.headers with its token. +runway_anton = production.override( + endpoints={ + 'ged112k': 'wss://runway-pythagoras-dev--curie-ged112k-curieserver-web.modal.run', + 'gd127k': 'wss://runway-pythagoras-dev--curie-gd127k-curieserver-web.modal.run', + 'gdf127k': 'wss://runway-pythagoras-dev--curie-gdf127k-curieserver-web.modal.run', + 'gdl142k': 'wss://runway-pythagoras-dev--curie-gdl142k-curieserver-web.modal.run', + }, + sampler=balanced, + group_fields=EVAL_GROUP_FIELDS, +) + +runway_ziyi = production.override( + endpoints={ + 'fm150k_us': 'wss://runway-pythagoras-dev--gyros-fm150k-us-gyrosserver-web.modal.run', + 'fm_tsu_actuni_150k': 'wss://runway-pythagoras-dev--gyros-fm-tsu-actuni-150k-gyrosserver-web.modal.run', + }, + sampler=balanced, + group_fields=EVAL_GROUP_FIELDS, ) diff --git a/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index 38c7bd501..03df6eed0 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -135,6 +135,13 @@ def main( """ assert (driver is None) != (evals is None), 'Provide exactly one of driver or evals' + # Ahead of warmup, which opens a session and so already samples: the counter it samples against holds + # what the output directory already recorded, rather than starting the run from zero. + if output_dir is not None: + output_dir = pos3.sync(output_dir, sync_on_error=True) + utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml']) + _seed_counter(policy, output_dir) + # Drive the policy's remote endpoints through their cold start before hardware and the operator # surface come up: opening a session blocks on the server handshake, which returns only once the # model is loaded, and a SampledPolicy reaches every sub-policy. The first episode then begins @@ -144,11 +151,6 @@ def main( logger.info('Warming up policy endpoints') policy.new_session().close() - if output_dir is not None: - output_dir = pos3.sync(output_dir, sync_on_error=True) - utils.save_run_metadata(output_dir, patterns=['*.py', '*.toml']) - _seed_counter(policy, output_dir) - # One completion sink — so one ``SampledPolicy`` counter — across every eval, keeping sampling balanced # over the whole sweep. on_complete = _completion_sink(policy) diff --git a/positronic/offboard/client.py b/positronic/offboard/client.py index 2c1c1ff7f..4fc248a9b 100644 --- a/positronic/offboard/client.py +++ b/positronic/offboard/client.py @@ -180,11 +180,12 @@ def new_session(self) -> InferenceSession: if ws is not None: ws.close() # A non-101 upgrade response only means "not ready" when it's a 5xx or 429; any other status - # (401/403/404, …) is permanent misconfiguration and surfaces immediately. + # (401/403/404, …) is permanent misconfiguration and surfaces immediately — naming the URL, + # which is the whole diagnosis when several endpoints are being opened at once. if isinstance(e, InvalidStatus) and not ( e.response.status_code >= 500 or e.response.status_code == 429 ): - raise + raise RuntimeError(f'{e} (connecting to {self.session_url})') from e if time.monotonic() >= deadline: raise TimeoutError(f'{e} (connecting to {self.session_url})') from e logger.info('Server not ready (cold start?): %s; retrying in %.0fs', e, backoff) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 55eb7c0c0..dc39b6b01 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -255,6 +255,20 @@ def test_remote_policy_meta_exposes_server_fields(): assert meta['server.model_name'] == 'foo' +def test_label_distinguishes_endpoints_serving_one_checkpoint(): + """Two deployments of the same checkpoint report identical server metadata, so their labels are + what a ``SampledPolicy`` keys on.""" + server_meta = {'checkpoint_path': '/ckpts/abc', **EMPTY_STACK} + labelled, _ = _mock_remote_policy(server_meta, label='nm167k') + other, _ = _mock_remote_policy(server_meta, label='nm167k_us') + unlabelled, _ = _mock_remote_policy(server_meta) + + assert labelled.meta['label'] == 'nm167k' + assert other.meta['label'] == 'nm167k_us' + assert labelled.meta['server.checkpoint_path'] == other.meta['server.checkpoint_path'] + assert 'label' not in unlabelled.meta + + def test_no_declaration_falls_back_to_chunked_schedule(): """A server that declares no ``local_stack`` in the handshake gets the standard ChunkedSchedule.""" clock = [0.0] diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 8e0f4012d..9497c8d6c 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,12 +1,17 @@ from __future__ import annotations +import logging from abc import ABC, abstractmethod from collections import Counter from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from operator import attrgetter from typing import Any from positronic.policy.sampler import EpisodeCounter, Sampler, UniformSampler +logger = logging.getLogger(__name__) + Now = Callable[[], float] # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. @@ -258,7 +263,12 @@ def sampler(self) -> Sampler | None: def _get_keys(self) -> tuple[str, ...]: if self._keys is None: - keys = tuple(p.meta.get(self._key_field, str(i)) for i, p in enumerate(self._policies)) + # A sub-policy's meta can cost a whole model load to read: a remote one holds a session open + # until its backend answers. Read them at once, so the first ``new_session`` waits out the + # slowest sub-policy rather than the sum of all of them. + with ThreadPoolExecutor(max_workers=len(self._policies)) as pool: + metas = list(pool.map(attrgetter('meta'), self._policies)) + keys = tuple(meta.get(self._key_field, str(i)) for i, meta in enumerate(metas)) duplicates = sorted(k for k, n in Counter(keys).items() if n > 1) if duplicates: raise ValueError( @@ -271,7 +281,11 @@ def _get_keys(self) -> tuple[str, ...]: def new_session(self, context=None, now=None): keys = self._get_keys() ctx = context or {} - key = self.sampler.sample(keys, ctx, self.counter.counts(keys, ctx)) + counts = self.counter.counts(keys, ctx) + key = self.sampler.sample(keys, ctx, counts) + # Which sub-policy ran is the same fact whichever strategy picked it, so it is reported here rather + # than per sampler. It names the episode's policy, so a blinded operator must not watch this log. + logger.info('Sampled %r; completed so far: %s', key, counts) policy = self._policies[keys.index(key)] session = policy.new_session(context, now) return _KeyedSession(session, policy.meta, self._key_field, key) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 4ddc38539..f5f32b831 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -156,12 +156,17 @@ class RemotePolicy(Policy): ``local`` and ``compress_images`` stand in for a server that declares neither — see ``_operator_override``. ``recording_dir`` taps the raw and wire boundaries around the stack. + + ``label`` is the operator's name for this endpoint, reported as its own meta field. Two deployments + of one checkpoint differ only in how they serve it, so the server's own metadata cannot tell them + apart — the operator's name is what distinguishes them, in sampling and in the recorded episodes. """ def __init__( self, url: str, *, + label: str | None = None, local: PolicyWrapper | None = None, recording_dir: str | None = None, headers: dict[str, str] | None = None, @@ -169,6 +174,7 @@ def __init__( compress_images: bool | None = None, ): self._endpoint = _Endpoint(url, headers=headers, infer_timeout=infer_timeout, compress_images=compress_images) + self._label = label self._local = local self._recording_dir = pos3.sync(recording_dir) if recording_dir else None self._stacked: Policy | None = None @@ -215,7 +221,8 @@ def new_session(self, context=None, now=None) -> Session: @property def meta(self) -> dict[str, Any]: - return self._policy().meta + meta = self._policy().meta + return meta if self._label is None else meta | {'label': self._label} def close(self): self._endpoint.close() diff --git a/positronic/policy/sampler.py b/positronic/policy/sampler.py index 46ed0afd5..4206f4b3b 100644 --- a/positronic/policy/sampler.py +++ b/positronic/policy/sampler.py @@ -1,12 +1,9 @@ -import logging import random from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Hashable, Sequence from typing import Any -logger = logging.getLogger(__name__) - class EpisodeCounter: """Per-key tally of completed episodes, grouped by context fields. @@ -99,10 +96,4 @@ def sample(self, keys: Sequence[str], context: dict[str, Any], counts: dict[str, c = [counts[k] for k in keys] max_count = max(c) if c else 0 weights = [max_count + self._balance - x for x in c] - chosen = random.choices(list(keys), weights)[0] - lines = ['BalancedSampler'] - for k, x, w in zip(keys, c, weights, strict=True): - marker = ' ←' if k == chosen else '' - lines.append(f' {k}: count={x} weight={w}{marker}') - logger.info('\n'.join(lines)) - return chosen + return random.choices(list(keys), weights)[0] diff --git a/positronic/policy/tests/test_sampler.py b/positronic/policy/tests/test_sampler.py index a5c4b7a3d..81f367b6e 100644 --- a/positronic/policy/tests/test_sampler.py +++ b/positronic/policy/tests/test_sampler.py @@ -1,3 +1,4 @@ +import threading from typing import Any from positronic.policy.base import Policy, SampledPolicy, Session @@ -32,6 +33,23 @@ def meta(self): return self._meta +class BarrierPolicy(StubPolicy): + """Stands in for a cold remote endpoint: the first ``meta`` read returns only once every peer is also + inside its own first read. Later reads are free, as they are once an endpoint has its server metadata.""" + + def __init__(self, barrier: threading.Barrier, meta: dict[str, Any]): + super().__init__(meta=meta) + self._barrier = barrier + self._cold = True + + @property + def meta(self): + if self._cold: + self._cold = False + self._barrier.wait() + return self._meta + + def _session(key, key_field='ckpt'): return _StubSession(meta={key_field: key}) @@ -158,6 +176,16 @@ def test_sampled_policy_discovers_keys_from_meta(): assert sampled._keys == ('/path/a', '/path/b') +def test_key_discovery_reads_every_sub_policy_meta_at_once(): + """Reading N cold sub-policies costs one model load, not N: each ``meta`` here blocks until all are in.""" + barrier = threading.Barrier(3, timeout=5) + sampled = SampledPolicy(*(BarrierPolicy(barrier, {'ckpt': k}) for k in 'abc'), key_field='ckpt') + + sampled.new_session({}) + + assert sampled._keys == ('a', 'b', 'c') + + def test_sampled_policy_delegates_to_sampler(): p1 = StubPolicy(meta={'ckpt': '/path/a'}) p2 = StubPolicy(meta={'ckpt': '/path/b'})