feat: improve typing throughout the codebase #39
Conversation
rayokamoto
force-pushed
the
feat/typing
branch
from
b85ab4d to
8d33c17
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughCentralizes domain/types in Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
kinitro/environments/genesis/scene_generator.py (1)
130-143:⚠️ Potential issue | 🔴 Critical
ObjectTypeenum stringifies as"ObjectType.BOX", not"box"— breaksobject_idformat.
ObjectTypeextendsEnum(notstr, Enum), sof"{obj_type}"produces"ObjectType.BOX"instead of"box". The generatedobject_idwill change from e.g."obj_00_red_box"to"obj_00_red_ObjectType.BOX", which may break downstream consumers that parse or match on the id.Use
obj_type.valuein the f-string:🐛 Proposed fix
objects.append( SceneObjectConfig( - object_id=f"obj_{i:02d}_{color_name}_{obj_type}", + object_id=f"obj_{i:02d}_{color_name}_{obj_type.value}", object_type=obj_type,kinitro/environments/registry.py (1)
106-116:⚠️ Potential issue | 🔴 CriticalUse
.valueto handleEnvironmentFamilyenum members correctly in f-string formatting.When
familyis anEnvironmentFamilyenum member,f"{family}/"produces"EnvironmentFamily.METAWORLD/"in Python 3.12 (not"metaworld/"), causing thestartswith()check to fail. This breaks the type contract since the function signature explicitly acceptsstr | EnvironmentFamily. Use.valueto ensure enum members are properly converted:Proposed fix
def get_environments_by_family(family: str | EnvironmentFamily) -> list[EnvironmentId]: - return [EnvironmentId(env_id) for env_id in ENVIRONMENTS if env_id.startswith(f"{family}/")] + family_prefix = family.value if isinstance(family, EnvironmentFamily) else family + return [EnvironmentId(env_id) for env_id in ENVIRONMENTS if env_id.startswith(f"{family_prefix}/")]
🤖 Fix all issues with AI agents
In `@kinitro/chain/commitments.py`:
- Line 322: committed_block can be None when passed into BlockNumber, which
silently allows None into an int-typed field; update the code that constructs
BlockNumber(committed_block) (from the _query_commitment_by_hotkey result or
neuron.last_update) to guard against None: check the returned committed_block
for None and either skip adding that commitment record or supply a safe
default/raise, e.g., only call BlockNumber(...) and use the commitment if
committed_block is not None, otherwise omit the commitment entry (or explicitly
handle the None case according to domain semantics).
In `@kinitro/environments/genesis/task_types.py`:
- Line 129: The error message uses the ObjectType enum directly so f-string
interpolation shows the enum repr (e.g., "ObjectType.BOX"); update the message
construction in the FeasibilityResult return(s) to use target.object_type.value
instead of target.object_type (and similarly for the other occurrence around
line 137) so the message reads the plain value (e.g., "box") — adjust the
f-strings where FeasibilityResult(False, f"Object {target.object_type} is not
{prop}") appears to use .value.
In `@kinitro/scoring/winners_take_all.py`:
- Around line 178-191: Hoist the conversion of subset_weight_scheme to a
SubsetWeightScheme once before the for loop (avoid re-creating `scheme` each
iteration) and move only the computation that depends on `subset_size` into the
loop; use the `SubsetWeightScheme` enum (or symbol `subset_weight_scheme`) to
resolve the scheme once, then inside the loop compute `subset_weight` based on
`subset_size`. Also make the match/case exhaustive by adding a wildcard `case
_:` that assigns a safe default or raises a clear error so `subset_weight` is
always defined (refer to the names `SubsetWeightScheme`, `subset_weight_scheme`,
`scheme`, `subset_size`, and `subset_weight` to locate the change).
In `@kinitro/types.py`:
- Around line 91-103: TaskCreateData currently uses TypedDict with total=False
which makes every key optional; change TaskCreateData to use the default
total=True and mark only optional fields with NotRequired (import NotRequired
from typing), i.e. keep required keys (cycle_id, miner_uid, miner_hotkey,
miner_endpoint, env_id, seed) as regular annotations and mark task_uuid,
miner_repo, miner_revision as NotRequired[...] so create_tasks_bulk (which
indexes required keys directly and only .get() for the three optional fields)
matches the type definition.
🧹 Nitpick comments (14)
kinitro/environments/genesis/envs/g1_humanoid.py (1)
61-92: Missing defaultcase _branch in_compute_reward.
_check_success(line 129) has acase _: return Falsefallback, but_compute_rewardsilently returns only the alive bonus for an unrecognizedTaskType. Adding a default case (even if it's just a pass or a log) would keep both methods consistent and make it easier to catch new task types that haven't been wired up yet.Suggested fix
case TaskType.PUSH: obj_pos = object_states.get(task_spec.target_object_id) dest_pos = task_spec.destination_position if obj_pos is not None and dest_pos is not None: dist_to_dest = float(np.linalg.norm(obj_pos[:2] - np.array(dest_pos[:2]))) reward += max(0.0, 1.0 - dist_to_dest / 5.0) * 0.1 + + case _: + pass # Unknown task type; only alive bonus applies return rewardkinitro/environments/metaworld_env.py (1)
584-603: Type mismatch: MetaWorld's info dict is assigned directly toStepInfo.On Line 590,
env.step(mw_action)returns MetaWorld's own info dict (a plaindict[str, Any]), which is unpacked intoinfo: StepInfo. While this works at runtime (TypedDicts are just dicts), MetaWorld's info may contain keys not declared inStepInfo. Since the env is already accessed viacast(Any, ...), the type checker won't catch this, but it's worth noting that the returnedStepInfomay carry extra undeclared keys.Not blocking — just a heads-up for future type-strictness efforts.
kinitro/executor/api_client.py (1)
39-43: Consider simplifying the union type.
list[str] | list[EnvironmentId]is redundant sinceEnvironmentIdis aNewTypeoverstr— alist[EnvironmentId]already satisfieslist[str]at runtime. You could simplify tolist[EnvironmentId] | Noneto encourage callers to adopt the new type.That said, keeping both during a transitional period is understandable.
environments/metaworld/env.py (1)
62-66: Consider typingenv_idparameters asEnvironmentIdfor consistency.
list_environmentsnow returnslist[EnvironmentId], butevaluate,_run_evaluation,_get_env, and_build_error_resultstill annotateenv_idasstr. SinceEnvironmentIdis aNewType(str), this is not a bug, but applying it here would be more consistent with the PR's goals.Also applies to: 99-111, 179-191, 317-325
kinitro/scoring/threshold.py (1)
43-43:episodes_per_envdict variant still typed asdict[str, int]rather thandict[EnvironmentId, int].Since
env_idkeys inminer_scoresare nowEnvironmentId, the dict form ofepisodes_per_envshould match for consistency.Suggested fix
def compute_miner_thresholds( miner_scores: MinerScores, - episodes_per_env: int | dict[str, int], + episodes_per_env: int | dict[EnvironmentId, int], z_score: float = 1.5, min_gap: float = 0.02, max_gap: float = 0.10, ) -> MinerThresholds:kinitro/rl_interface.py (1)
153-154: Consider narrowing thelisttype in the union.
EncodedImage | listleaves the list element type unspecified. If the legacy format is always nested numeric lists, a narrower annotation (e.g.,list[list[list[int]]]or at minimumlist[Any]) would be more self-documenting, though this is a minor nit given the backward-compat intent.kinitro/scheduler/main.py (1)
257-257: RedundantEnvironmentId()wrapping —self.env_idsis alreadylist[EnvironmentId].
self.env_idsis assigned aslist[EnvironmentId](line 44) or fromget_all_environment_ids()which also returnslist[EnvironmentId](line 66). The comprehension[EnvironmentId(e) for e in self.env_ids]is a no-op re-wrap.♻️ Simplify
weights, weights_u16 = compute_weights( miner_scores=miner_scores, - env_ids=[EnvironmentId(e) for e in self.env_ids], + env_ids=self.env_ids, episodes_per_env=self.config.episodes_per_env,kinitro/api/routes/miners.py (1)
61-63: Minor:env_idsis alreadylist[EnvironmentId], so the wrapping in the comprehension is a no-op.
get_all_environment_ids()returnslist[EnvironmentId], makingEnvironmentId(env_id)on line 62 (and again on lines 79, 84) redundant. Harmless at runtime sinceNewTypeis transparent, but it adds noise.kinitro/scoring/pareto.py (1)
102-106:n_samples_per_envdict key type inconsistent withEnvironmentId.
env_idsis nowlist[EnvironmentId], butn_samples_per_envacceptsdict[str, int]. For type-checker consistency, this should bedict[EnvironmentId, int]to match the keys used insamples.get(env_ids[j], 50)on line 152.Proposed fix
def compute_pareto_frontier( miner_scores: MinerScores, env_ids: list[EnvironmentId], - n_samples_per_env: int | dict[str, int], + n_samples_per_env: int | dict[EnvironmentId, int], ) -> ParetoResult:kinitro/scheduler/scoring.py (1)
171-175:miners_by_uidvalue type should beHotkeyinstead ofstr.The comment says
uid -> hotkey, and the values are used to populate thehotkeyfield ofMinerScoreData(which expectsHotkey). For consistency with the broader typing effort:Proposed fix
def convert_to_scores_data( miner_scores: MinerScores, - miners_by_uid: dict[MinerUID, str], # uid -> hotkey + miners_by_uid: dict[MinerUID, Hotkey], # uid -> hotkey episodes_per_env: int, ) -> list[MinerScoreData]:This would also require adding
Hotkeyto the import on line 14.kinitro/environments/genesis/base.py (1)
751-751: Narrow the local type annotations to match actual usage.The
rgbanddepthdicts in_build_observationonly storeEncodedImagevalues fromencode_image()calls; the| listpart is unnecessary here. While theObservationclass itself defines these fields with| list, the local variables can safely use a narrower type that reflects what's actually stored.Proposed simplification
- rgb: dict[str, EncodedImage | list] = {} + rgb: dict[str, EncodedImage] = {}- depth: dict[str, EncodedImage | list] = {} + depth: dict[str, EncodedImage] = {}kinitro/types.py (2)
1-14: Well-structured centralized type module.Clean separation with zero intra-project imports, clear docstring, and proper import grouping. One minor note: importing
numpyat the module level means every consumer of even simple aliases likeMinerUIDorEnvironmentIdtransitively imports NumPy. If startup time or lightweight consumers ever become a concern, the NumPy-dependent TypedDicts (RobotStateDict,ProceduralTaskResult) could be split into a separate sub-module.
155-161:dict[str, object]is overly restrictive for a randomization bag.
objectis the top type but disallows attribute access and most operations without casting. If these values are truly heterogeneous (floats, lists, bools, etc.),Anyis the pragmatic choice here — or afloat | bool | list[float]union if the set of possible types is known. As per coding guidelines,Anyshould include a comment explaining why.kinitro/scoring/winners_take_all.py (1)
152-153: Consider using the enum member as the default instead of a raw string.The default
"linear"works because of theisinstance+ conversion logic, but usingSubsetWeightScheme.LINEARdirectly communicates intent more clearly and avoids the conversion path when the default is used.Proposed fix
- subset_weight_scheme: str | SubsetWeightScheme = "linear", + subset_weight_scheme: str | SubsetWeightScheme = SubsetWeightScheme.LINEAR,
rayokamoto
force-pushed
the
feat/typing
branch
from
8d33c17 to
0e339f6
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
kinitro/environments/genesis/scene_generator.py (1)
130-151:⚠️ Potential issue | 🟠 MajorBug:
ObjectTypeenum renders as"ObjectType.BOX"in the f-string, changingobject_idformat.
ObjectTypeextendsEnum(notstr, Enum), sof"{obj_type}"produces"ObjectType.BOX"instead of the previous"box". This changes the generatedobject_id(e.g.,"obj_00_red_ObjectType.BOX"instead of"obj_00_red_box"), which may break downstream logic that parses or matches these IDs.🐛 Proposed fix: use `.value` for the string representation
objects.append( SceneObjectConfig( - object_id=f"obj_{i:02d}_{color_name}_{obj_type}", + object_id=f"obj_{i:02d}_{color_name}_{obj_type.value}", object_type=obj_type,
🤖 Fix all issues with AI agents
In `@kinitro/chain/commitments.py`:
- Around line 466-468: The code assumes subtensor.set_commitment() returns an
object with a .success attribute but in bittensor 8.0.0 it returns a boolean;
update the handling in set_commitment (look for the variable result and the line
success = result.success) to treat result as a bool (e.g., set success = result)
and remove any subsequent attribute access on result so the function correctly
processes the boolean return value.
In `@kinitro/chain/weights.py`:
- Around line 81-95: subtensor.set_weights returns a tuple (success, message)
not an object, so change the code that assigns and checks result to unpack the
tuple (e.g., success, message = subtensor.set_weights(...)) instead of accessing
result.success; then use success for the if/return and include message in the
failure log (e.g., logger.error("weights_set_failed", message=message)) so the
actual error message surfaces rather than being swallowed.
In `@kinitro/cli/crypto_commands.py`:
- Line 8: Update the import so Wallet is imported from the separate package and
Subtensor remains from bittensor: replace the combined import of Wallet and
Subtensor with an import that gets Wallet from the bittensor_wallet package and
Subtensor from bittensor (or use bittensor as bt and access bt.Subtensor);
update any references in this module to use the new import if namespaced
differently; and add bittensor-wallet to the project dependencies (e.g.,
requirements or pyproject) so Wallet is available at runtime.
- Line 242: The code assumes subtensor.set_commitment() returns an object with a
.success attribute but in bittensor 8.0.0 it returns a boolean; update the
assignment where result is set by subtensor.set_commitment(...) so that you use
the boolean directly (replace accessing result.success with success = result) —
adjust the handling around the result variable in the function that calls
subtensor.set_commitment to remove the .success access and treat result as a
bool.
In `@kinitro/cli/miner/deploy.py`:
- Line 7: The import line using "from bittensor import Subtensor, Wallet" is
incompatible with bittensor v8.0.0; replace it by importing bittensor as bt and
importing Wallet from the separate bittensor_wallet package, i.e., use "import
bittensor as bt" and "from bittensor_wallet import Wallet", then update any
direct references to Subtensor to bt.Subtensor (e.g., where Subtensor(...) is
used) and ensure the project dependencies include bittensor-wallet so the new
Wallet import resolves.
In `@kinitro/validator/main.py`:
- Around line 42-46: Validator.__init__ currently instantiates Subtensor
synchronously (self.subtensor = Subtensor(...)), which blocks the event loop
because run_validator calls the constructor from async code; replace this by
deferring the blocking work: either instantiate AsyncSubtensor and call its
awaitable initialize() from run_validator (or an async factory on Validator), or
move the Subtensor(...) call onto a worker thread with asyncio.to_thread and
assign the result to self.subtensor so the event loop is not blocked; update
Validator.__init__/run_validator to use the chosen approach and ensure any
subsequent uses of self.subtensor await initialize() or run after the to_thread
completion.
🧹 Nitpick comments (12)
kinitro/environments/genesis/envs/g1_humanoid.py (1)
61-91: Missing default case in_compute_rewardmatch-case.
_check_success(line 129) includes acase _: return Falsedefault branch, but_compute_rewardhas no corresponding default. If a newTaskTypeis added to the enum, this match will silently skip all task-specific reward terms without any signal. Adding a default here for symmetry — even if it just passes — improves clarity and helps catch future omissions.Suggested fix
case TaskType.PUSH: obj_pos = object_states.get(task_spec.target_object_id) dest_pos = task_spec.destination_position if obj_pos is not None and dest_pos is not None: dist_to_dest = float(np.linalg.norm(obj_pos[:2] - np.array(dest_pos[:2]))) reward += max(0.0, 1.0 - dist_to_dest / 5.0) * 0.1 + + case _: + pass # Only alive bonus for unknown task types return rewardenvironments/genesis/env.py (1)
103-105: Consider tightening theAffinetesEnvprotocol return type.The
AffinetesEnvprotocol inkinitro/types.py(line 172) declareslist_environmentsas-> Any, but all concrete implementations (genesis/env.py,metaworld/env.py, template) now returnlist[EnvironmentId]. Updating the protocol to match would provide better type safety.kinitro/executor/api_client.py (1)
42-42: Consider simplifying the union type forenv_ids.
list[str] | list[EnvironmentId]prevents mixed lists. If callers might pass a heterogeneous list,list[str | EnvironmentId] | Noneis more permissive. If the intent is to eventually migrate all callers toEnvironmentId, consider justlist[EnvironmentId] | None— sinceEnvironmentIdis aNewTypeoverstr, callers can wrap withEnvironmentId(...)at the boundary.environments/metaworld/env.py (1)
343-345: Missing-> Nonereturn type oncleanup.Public method is missing its return type annotation. As per coding guidelines, "Add return types to public functions and methods."
Proposed fix
- async def cleanup(self): + async def cleanup(self) -> None:kinitro/scheduler/task_generator.py (1)
95-107: Task dict values are not wrapped with the new typed constructors, unlike every other file in this PR.The return type is
list[TaskCreateData], which expectsMinerUID,Hotkey,Seed,EnvironmentId, andTaskUUIDtyped values. However, the dict is constructed with raw values (miner.uid,miner.hotkey,seedasint, etc.) without wrapping them (e.g.,MinerUID(miner.uid),Seed(seed),TaskUUID(task_uuid)).While this works at runtime since
NewTypeis erased, it defeats the purpose of this PR's type-safety improvements and a type checker likemypy --strictorpyrightwill flag these as errors.♻️ Proposed fix: wrap values with typed constructors
You'd need to add
Hotkey,MinerUID,Seed,TaskUUIDto the import on line 9, then:tasks.append( { - "task_uuid": task_uuid, + "task_uuid": TaskUUID(task_uuid), "cycle_id": cycle_id, - "miner_uid": miner.uid, - "miner_hotkey": miner.hotkey, + "miner_uid": MinerUID(miner.uid), + "miner_hotkey": Hotkey(miner.hotkey), "miner_endpoint": endpoint, "miner_repo": miner.huggingface_repo, "miner_revision": miner.revision_sha, "env_id": env_id, - "seed": seed, + "seed": Seed(seed), } )kinitro/api/routes/miners.py (1)
61-84: RedundantEnvironmentId()wrapping —env_idsare alreadyEnvironmentId.
get_all_environment_ids()already returnslist[EnvironmentId], soenv_idin the loop is already typed asEnvironmentId. The re-wrapping on lines 62, 79, and 84 is unnecessary.♻️ Remove redundant wrapping
env_stats: dict[EnvironmentId, EnvStatsEntry] = { - EnvironmentId(env_id): EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids + env_id: EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids }- stats = env_stats[EnvironmentId(env_id)] + stats = env_stats[env_id]result.append( EnvironmentInfo( - env_id=EnvironmentId(env_id), + env_id=env_id,kinitro/scoring/pareto.py (1)
102-106:n_samples_per_envdict variant still usesstrkeys instead ofEnvironmentId.The
env_idsparameter is nowlist[EnvironmentId], and on line 152 the code doessamples.get(env_ids[j], 50)whereenv_ids[j]isEnvironmentId. The dict type should match for type-checker consistency.♻️ Align the dict key type
def compute_pareto_frontier( miner_scores: MinerScores, env_ids: list[EnvironmentId], - n_samples_per_env: int | dict[str, int], + n_samples_per_env: int | dict[EnvironmentId, int], ) -> ParetoResult:kinitro/scheduler/main.py (1)
248-248: RedundantEnvironmentIdwrapping —self.env_idsis alreadylist[EnvironmentId].
self.env_idsis typed and populated aslist[EnvironmentId](lines 44, 53, 66), so re-wrapping each element is a no-op at runtime (NewType is the identity function). You can passself.env_idsdirectly.Suggested simplification
weights, weights_u16 = compute_weights( miner_scores=miner_scores, - env_ids=[EnvironmentId(e) for e in self.env_ids], + env_ids=self.env_ids, episodes_per_env=self.config.episodes_per_env,kinitro/backend/storage.py (1)
264-282:weightsdict hasMinerUID(int) keys but the ORM column is typeddict[str, float].
save_weightsacceptsdict[MinerUID, float](line 269), butComputedWeightsORM.weights_jsonis annotated asMapped[dict[str, float]](models.py line 124). JSONB serialization will stringify the int keys, so this works at runtime, but the type mismatch may confuse static analysis and readers.Consider either:
- Explicitly converting keys to strings before storing, or
- Updating the ORM annotation to
dict[int, float]to match the semantic intent.kinitro/scheduler/scoring.py (1)
171-205: Return type annotation sayslist[MinerScoreData]but the function appends plaindictliterals.This works at runtime because
TypedDictis structurally typed, but static type checkers (mypy/pyright) will flag the plain dict literal as incompatible withMinerScoreDataunless you explicitly annotate or construct it. Consider annotating the dict literal to help the checker:entry: MinerScoreData = { ... } scores_data.append(entry)This is minor since it's already correct at runtime.
kinitro/types.py (2)
164-175:Anyusage inAffinetesEnvProtocol is acceptable but worth a brief inline comment.The
Anyannotations onlist_environments,evaluate, and**kwargsare justified since this is a structural interface for an external package (affinetes) not available at type-check time. A short# External package – exact types unknownnote would satisfy the coding guideline to explainAnyusage. As per coding guidelines: "AvoidAnytype annotations unless required and explain why in code comments."
155-162:dict[str, object]fordomain_randomization— intentional or should this beAny?
objectis more restrictive thanAny(no attribute access or method calls without casting). If consumers need to access values without type narrowing,Anywith a comment may be more practical. If the intent is to force explicit narrowing,objectis correct.
rayokamoto
force-pushed
the
feat/typing
branch
2 times, most recently
from
7ed57b6 to
8b42a0e
Compare
|
@coderabbitai please re-review |
|
✅ Actions performedFull review triggered. |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
kinitro/environments/genesis/scene_generator.py (1)
130-143:⚠️ Potential issue | 🟠 MajorUse
obj_type.valuein the f-string to preserve the object_id format.
ObjectTypeis a plainEnum, not aStrEnum. When you usef"{obj_type}"whereobj_type = ObjectType.BOX, Python's__str__()returns"ObjectType.BOX"rather than"box". This changes the object_id format from the expectedobj_00_red_boxtoobj_00_red_ObjectType.BOX.Test fixtures hardcode the old format (
object_id="obj_00_red_box","obj_01_green_sphere", etc.) and will fail with the current code. Useobj_type.valueto preserve the original format:Proposed fix
objects.append( SceneObjectConfig( - object_id=f"obj_{i:02d}_{color_name}_{obj_type}", + object_id=f"obj_{i:02d}_{color_name}_{obj_type.value}", object_type=obj_type,kinitro/backend/storage.py (1)
264-282:⚠️ Potential issue | 🟡 Minor
weightsdict keys areMinerUID(int) but the JSONB column is typeddict[str, float].
ComputedWeightsORM.weights_jsonisMapped[dict[str, float]](models.py line 124), butweightshere isdict[MinerUID, float]whereMinerUIDisNewType(..., int). JSON serialization will stringify the keys, so it works at runtime, but static type checkers will flag the assignmentweights_json=weights. Consider converting keys explicitly:Suggested fix
weights_orm = ComputedWeightsORM( cycle_id=cycle_id, block_number=block_number, - weights_json=weights, + weights_json={str(k): v for k, v in weights.items()}, weights_u16_json=weights_u16, created_at=datetime.now(timezone.utc), )
🤖 Fix all issues with AI agents
In `@AGENTS.md`:
- Line 121: Update the American English spelling by changing the phrase
"Centralise shared newtypes, type aliases, enums, and `TypedDict` definitions in
`kinitro/types.py`. Import from there rather than re-defining types locally." to
use "Centralize" instead of "Centralise" in AGENTS.md so it matches the rest of
the document's spelling.
In `@pyproject.toml`:
- Line 11: The dependency line currently lists "bittensor>=10.1.0", which is
unsatisfiable because 10.1.0 doesn't exist on PyPI; update the requirement in
pyproject.toml where the string "bittensor>=10.1.0" appears to a valid
constraint such as "bittensor>=10.0.1,<11.0.0" (or pin to "bittensor==10.0.1")
so installations succeed.
- Line 30: Replace the unsatisfiable dependency "typer>=0.23.0" in
pyproject.toml with a real Typer release; change the requirement to a valid
version such as "typer>=0.22.0" or pin to "typer==0.22.0" (or use a bounded
range like "typer>=0.22.0,<0.23") so dependency resolution succeeds — update the
string "typer>=0.23.0" accordingly.
🧹 Nitpick comments (17)
pyproject.toml (1)
59-65: Duplicate dev dependency sections with diverging versions.Both
[project.optional-dependencies].dev(Lines 59–65) and[dependency-groups].dev(Lines 115–121) declare dev dependencies, but with mismatched version bounds:
Package optional-dependenciesdependency-groupspytest >=9.0.0>=9.0.2pytest-asyncio >=0.21.0>=1.3.0ruff >=0.1.0>=0.14.13pytest-cov >=4.0.0(missing) ty >=0.0.14>=0.0.14This will cause confusion about which versions are actually required. Consider consolidating into a single source of truth — likely
[dependency-groups]if your tooling supports PEP 735, or[project.optional-dependencies]otherwise.Also applies to: 115-121
kinitro/scoring/threshold.py (1)
43-43: Consider typingepisodes_per_envdict key asEnvironmentIdfor consistency.Since
miner_scoresnow usesEnvironmentIdkeys (viaMinerScores), thedict[str, int]variant ofepisodes_per_envcould bedict[EnvironmentId, int]to keep the domain typing consistent throughout this function's signature.Proposed fix
def compute_miner_thresholds( miner_scores: MinerScores, - episodes_per_env: int | dict[str, int], + episodes_per_env: int | dict[EnvironmentId, int], z_score: float = 1.5, min_gap: float = 0.02, max_gap: float = 0.10, ) -> MinerThresholds:kinitro/environments/genesis/envs/g1_humanoid.py (1)
61-92: Minor:_compute_rewardlacks acase _branch, unlike_check_success.If a new
TaskTypevariant is added,_compute_rewardsilently returns only the alive bonus (0.01) without any indication, while_check_successexplicitly handles the unknown case withcase _: return False. Adding a symmetriccase _(even if just a pass/comment) would make the intent clear and keep the two methods consistent.Proposed fix
case TaskType.PUSH: obj_pos = object_states.get(task_spec.target_object_id) dest_pos = task_spec.destination_position if obj_pos is not None and dest_pos is not None: dist_to_dest = float(np.linalg.norm(obj_pos[:2] - np.array(dest_pos[:2]))) reward += max(0.0, 1.0 - dist_to_dest / 5.0) * 0.1 + + case _: + pass # Unknown task type — alive bonus only return rewardenvironments/_template/env.py (1)
36-44: Template still has several untyped signatures — consider aligning with the genesis counterpart.Since this PR is about improving typing, the template (
__init__,_get_env,_build_error_result,evaluate) could benefit from the same treatment applied toenvironments/genesis/env.py(e.g.,-> None,-> RoboticsEnvironment,-> EvalResult). Template code tends to be copied verbatim into new environments, so stronger types here pay forward.kinitro/types.py (2)
164-175:AffinetesEnv.list_environmentsandevaluateuseAnywithout the required justification comment.Per the coding guidelines and AGENTS.md Line 119: "avoid
Anyunless required and explain why in code." All concrete implementations returnlist[EnvironmentId]anddict[str, Any]respectively, so at minimumlist_environmentscould be tightened. IfAnyis intentional (e.g., to keep the protocol loose for third-party environments), a brief comment explaining why would satisfy the guideline.Proposed fix (tighten where possible, comment where not)
class AffinetesEnv(Protocol): """Structural interface for affinetes-managed evaluation environments. Used by the executor subsystem to interact with Docker/Basilica environments without depending on the ``affinetes`` package at type-checking time. """ def is_ready(self) -> bool: ... - async def list_environments(self) -> Any: ... - async def evaluate(self, **kwargs: Any) -> dict[str, Any]: ... + async def list_environments(self) -> list[EnvironmentId]: ... + # evaluate() uses Any because affinetes actors accept heterogeneous kwargs + async def evaluate(self, **kwargs: Any) -> dict[str, Any]: ... async def cleanup(self) -> None: ...As per coding guidelines: "Avoid
Anytype annotations unless required and explain why in code."
23-25: Consider usingTypeAlias(or the Python 3.12typestatement) for the composite aliases.Currently
MinerScores,MinerThresholds, andMinerFirstBlocksare bare assignments. Type checkers handle these, but usingTypeAlias(ortype MinerScores = ...in 3.12) makes the intent explicit and avoids potential confusion with runtime values.Example using Python 3.12 `type` statement
-MinerScores = dict[MinerUID, dict[EnvironmentId, float]] # uid -> env_id -> score -MinerThresholds = dict[MinerUID, dict[EnvironmentId, float]] # uid -> env_id -> threshold -MinerFirstBlocks = dict[MinerUID, BlockNumber] # uid -> block_number +type MinerScores = dict[MinerUID, dict[EnvironmentId, float]] # uid -> env_id -> score +type MinerThresholds = dict[MinerUID, dict[EnvironmentId, float]] # uid -> env_id -> threshold +type MinerFirstBlocks = dict[MinerUID, BlockNumber] # uid -> block_numberNote: The
typestatement requires removingfrom __future__ import annotationsor using it only in Python 3.12+. If you stay withfrom __future__ import annotations, useTypeAliasfromtypinginstead.kinitro/executor/api_client.py (1)
42-42: Consider simplifying the union type.
list[str] | list[EnvironmentId] | Nonecould be simplified tolist[EnvironmentId] | Noneif all callers have been migrated to useEnvironmentId. If backward compatibility with plain strings is needed,list[str | EnvironmentId] | Noneis more precise than two separate list types.kinitro/environments/metaworld_env.py (1)
584-590: Minor type imprecision: MetaWorld'senv.step()returns a broader dict thanStepInfo.On line 590,
infois reassigned from MetaWorld's nativeenv.step()which returns its own info dict containing keys not defined inStepInfo(e.g.,near_object,grasp_success, etc.). TheStepInfoannotation on line 584 is immediately overwritten and the actual runtime value won't strictly conform toStepInfo.This is pragmatically fine since
TypedDictwithtotal=Falsedoesn't enforce at runtime, but a# type: ignore[assignment]or a cast would make the intent explicit for type checkers.kinitro/environments/registry.py (1)
78-78:get_environmentparameter remainsstrwhile sibling functions now useEnvironmentId.Minor inconsistency:
get_all_environment_ids()andget_environments_by_family()returnEnvironmentId, butget_environment()accepts plainstr. SinceEnvironmentIdis aNewType(str), this works at runtime, but updating the signature toenv_id: EnvironmentIdwould make the public API consistent.kinitro/scoring/pareto.py (1)
102-106:n_samples_per_envdict key type is stillstrinstead ofEnvironmentId.
env_idsis nowlist[EnvironmentId]andminer_scoresusesEnvironmentIdkeys, butn_samples_per_envon line 105 isint | dict[str, int]. For consistency (and correct type-checking when used withsamples.get(env_ids[j], 50)on line 152), the dict variant should useEnvironmentIdkeys.Suggested fix
def compute_pareto_frontier( miner_scores: MinerScores, env_ids: list[EnvironmentId], - n_samples_per_env: int | dict[str, int], + n_samples_per_env: int | dict[EnvironmentId, int], ) -> ParetoResult:kinitro/environments/genesis/base.py (1)
744-775:_build_observationtypes are mostly clean, butdict[str, EncodedImage | list]is imprecise.Lines 751 and 755 use
dict[str, EncodedImage | list]forrgbanddepth. The| listarm is vague — if this is for the case where theObservationmodel expects an empty list as default, it would be clearer to let the dict simply be empty. If theObservationmodel genuinely requireslistas an alternative value type, adding a brief comment would help clarify intent.kinitro/scheduler/task_generator.py (1)
89-107: Task dict values not wrapped with corresponding NewTypes.The return type is
list[TaskCreateData], whereTaskCreateDataexpectsseed: Seedandtask_uuid: TaskUUID(both NewTypes). However,generate_seedreturns a plainintandstr(uuid.uuid4())returns a plainstr. For consistency with the typing effort in this PR, wrap these values:Suggested fix
+from kinitro.types import BlockNumber, EnvironmentId, Seed, TaskCreateData, TaskUUID ... - task_uuid = str(uuid.uuid4()) + task_uuid = TaskUUID(str(uuid.uuid4())) ... - seed = generate_seed(task_uuid) + seed = Seed(generate_seed(task_uuid))kinitro/scheduler/main.py (1)
248-248: RedundantEnvironmentIdwrapping.
self.env_idsis alreadylist[EnvironmentId](assigned at line 53 or 66), so[EnvironmentId(e) for e in self.env_ids]is a no-op copy. Passself.env_idsdirectly.Suggested fix
- env_ids=[EnvironmentId(e) for e in self.env_ids], + env_ids=self.env_ids,kinitro/api/routes/miners.py (2)
61-63: RedundantEnvironmentIdwrapping — values are alreadyEnvironmentId.
get_all_environment_ids()returnslist[EnvironmentId], soEnvironmentId(env_id)in the dict comprehension is unnecessary. Same applies on lines 79 and 84.Suggested fix
env_stats: dict[EnvironmentId, EnvStatsEntry] = { - EnvironmentId(env_id): EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids + env_id: EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids }- stats = env_stats[EnvironmentId(env_id)] + stats = env_stats[env_id]- env_id=EnvironmentId(env_id), + env_id=env_id,
29-46: O(n²) score aggregation per miner.Lines 42-45 re-scan the full
scoreslist for each miner to compute the average. Consider accumulating totals during the initial loop to avoid the quadratic cost as the number of miners/scores grows.kinitro/scheduler/scoring.py (2)
171-175:miners_by_uidvalue type should beHotkey, notstr.The parameter is typed
dict[MinerUID, str]butMinerScoreDataexpectshotkey: Hotkey. The caller (main.pyline 258) already passesm.hotkeywhich isHotkey. Update the annotation and the fallback:Suggested fix
def convert_to_scores_data( miner_scores: MinerScores, - miners_by_uid: dict[MinerUID, str], # uid -> hotkey + miners_by_uid: dict[MinerUID, Hotkey], # uid -> hotkey episodes_per_env: int, ) -> list[MinerScoreData]:- hotkey = miners_by_uid.get(uid, "unknown") + hotkey = miners_by_uid.get(uid, Hotkey("unknown"))
198-204: TODO comments formean_rewardandepisodes_completed/episodes_failed.These are hardcoded to
0.0/episodes_per_env/0. Would you like me to generate an implementation that computes these from the actual task results, or open an issue to track it?
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
kinitro/environments/registry.py (1)
106-116:⚠️ Potential issue | 🟡 MinorStale docstring: "currently only 'metaworld'" — Genesis is also a registered family.
Line 111 references only
'metaworld'but the registry also containsgenesis/g1-v0.Suggested fix
- family: Environment family (currently only 'metaworld') + family: Environment family (e.g., 'metaworld', 'genesis')kinitro/environments/genesis/scene_generator.py (1)
141-151:⚠️ Potential issue | 🔴 CriticalBug:
obj_typein f-string producesObjectType.BOXinstead ofbox.
str(ObjectType.BOX)yields"ObjectType.BOX", not"box". This changesobject_idfrom e.g."obj_00_red_box"to"obj_00_red_ObjectType.BOX", which will break any downstream code that parses or matches object IDs.Use
.valueto preserve the original string representation:🐛 Proposed fix
objects.append( SceneObjectConfig( - object_id=f"obj_{i:02d}_{color_name}_{obj_type}", + object_id=f"obj_{i:02d}_{color_name}_{obj_type.value}", object_type=obj_type,kinitro/scheduler/task_generator.py (1)
95-107:⚠️ Potential issue | 🟡 MinorOnly
task_uuidandseedneed to be wrapped with their domain-type constructors;miner_uidandminer_hotkeyare already correctly typed.The return type is
list[TaskCreateData], but the dict construction uses raw primitives fortask_uuid(astr) andseed(anint). However,miner.uidandminer.hotkeyare already typed asMinerUIDandHotkeyrespectively in theMinerCommitmentclass, so they do not need wrapping.♻️ Suggested fix
+from kinitro.types import BlockNumber, EnvironmentId, Seed, TaskCreateData, TaskUUID ... tasks.append( { - "task_uuid": task_uuid, + "task_uuid": TaskUUID(task_uuid), "cycle_id": cycle_id, "miner_uid": miner.uid, "miner_hotkey": miner.hotkey, "miner_endpoint": endpoint, "miner_repo": miner.huggingface_repo, "miner_revision": miner.revision_sha, "env_id": env_id, - "seed": seed, + "seed": Seed(seed), } )
🤖 Fix all issues with AI agents
In `@kinitro/scoring/pareto.py`:
- Around line 102-106: The type annotation for the function
compute_pareto_frontier is inconsistent: n_samples_per_env is typed as dict[str,
int] but is used with env_ids (EnvironmentId), e.g. samples.get(env_ids[j], 50);
update the annotation to use dict[EnvironmentId, int] (or Mapping[EnvironmentId,
int]) so n_samples_per_env, the local samples variable, and any related lookups
match EnvironmentId keys and satisfy type checkers.
In `@pyproject.toml`:
- Line 24: The inline comment on the "numpy>=2.0.1,<2.3" dependency incorrectly
references a non-existent bittensor version; update the comment to reference the
correct bittensor version (the actual version declared in pyproject.toml
elsewhere) and make the comment accurate and concise (e.g., replace "lower bound
from bittensor >=10.1" with the real version constraint that provides the lower
bound), leaving the dependency spec itself unchanged.
🧹 Nitpick comments (18)
tests/unit/test_crypto.py (1)
34-34: Consider adding a type annotation for theuuid_strparameter.Since this PR is about improving typing, the parametrized argument could also be annotated for consistency.
- def test_uuid_to_bytes(self, uuid_str) -> None: + def test_uuid_to_bytes(self, uuid_str: str) -> None:kinitro/environments/registry.py (1)
78-98:get_environmentstill acceptsstrrather thanEnvironmentId.For consistency with the typed return values of
get_all_environment_idsandget_environments_by_family, consider updating theenv_idparameter to acceptEnvironmentId. SinceEnvironmentIdis aNewTypeoverstr, this is a zero-cost annotation change that improves type-checker coverage.kinitro/cli/miner/commitment.py (1)
4-5: Minor inconsistency:from bittensor import Subtensorhere vsimport bittensor as btindeploy.py.Both work, but the codebase now uses two different import styles for the same package across CLI modules. Consider aligning on one style for consistency.
kinitro/executor/api_client.py (1)
42-42: The unionlist[str] | list[EnvironmentId]is redundant ifEnvironmentIdis aNewTypeofstr.Since
EnvironmentId = NewType("EnvironmentId", str), alist[EnvironmentId]is already assignable wherelist[str]is expected at runtime (NewType doesn't create a distinct runtime type). The union was likely added for backward compatibility with existing callers, which is fine, but consider simplifying tolist[EnvironmentId] | Noneto encourage callers to adopt the stronger type.kinitro/environments/metaworld_env.py (1)
515-518:StepInforeturn type is a reasonable approximation but not fully accurate for MetaWorld's info dict.MetaWorld's
env.step()returns an info dict with extra keys beyond whatStepInfodeclares (e.g.,near_object,grasp_success,grasp_reward). SinceStepInfousestotal=Falseand the runtime dict is a plaindict, this works at runtime and the key callers actually use (success) is declared inStepInfo. Just be aware that strict type checkers (pyright in strict mode) may flag the extra keys if the return type is narrowed elsewhere.kinitro/scheduler/task_generator.py (1)
14-32:generate_seedreturns a plainint, not aSeed.Since the PR introduces
Seedas a NewType, consider updating the return type ofgenerate_seedtoSeedfor consistency.♻️ Suggested fix
-def generate_seed(task_uuid: str) -> int: +def generate_seed(task_uuid: str) -> Seed: ... - return int.from_bytes(hash_bytes, byteorder="big") & 0x7FFFFFFF + return Seed(int.from_bytes(hash_bytes, byteorder="big") & 0x7FFFFFFF)kinitro/environments/genesis/envs/g1_humanoid.py (1)
61-92: Consider adding a wildcard case to_compute_rewardfor symmetry with_check_success.
_check_success(line 129) has acase _: return Falsewildcard, but_compute_rewardsilently falls through for unknown task types, returning only the alive bonus. Adding a wildcard (even if it justpasses or logs a warning) would make intent explicit and prevent silent bugs when newTaskTypevariants are added.kinitro/validator/main.py (1)
70-73: Sync RPC calls (verify_weight_setting_eligibility,set_weights) still block the event loop.Both functions perform synchronous network I/O (chain RPC) and are called directly from
asyncmethods. Consider wrapping them withasyncio.to_thread(...)for consistency with theSubtensorinitialization fix. This is a pre-existing concern, not introduced by this PR.As per coding guidelines, "Use
async deffor I/O and DB operations."Also applies to: 129-134
kinitro/rl_interface.py (1)
154-155: Consider narrowing the legacylisttype annotation.
EncodedImage | listuses a barelistwhich is very broad. If the legacy format is always a nested list of integers (pixel data), considerlist[list[list[int]]]or at minimumlist[Any]to be explicit. Low priority since this is a transitional type for backward compatibility.kinitro/api/routes/miners.py (1)
61-63: RedundantEnvironmentId()wrapping —env_idsalready containsEnvironmentIdvalues.
get_all_environment_ids()returnslist[EnvironmentId](seekinitro/environments/registry.py, lines 100-102), soenv_idin the loop is already anEnvironmentId. The extraEnvironmentId(env_id)calls on lines 62, 79, and 84 are no-ops at runtime but add noise.♻️ Suggested cleanup
env_stats: dict[EnvironmentId, EnvStatsEntry] = { - EnvironmentId(env_id): EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids + env_id: EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids } ... - stats = env_stats[EnvironmentId(env_id)] + stats = env_stats[env_id] ... EnvironmentInfo( - env_id=EnvironmentId(env_id), + env_id=env_id,Also applies to: 79-79, 84-84
kinitro/scheduler/main.py (2)
248-248: RedundantEnvironmentIdre-wrapping.
self.env_idsis alreadylist[EnvironmentId](set at line 53 or 66), so[EnvironmentId(e) for e in self.env_ids]creates a new list of identical values. You can passself.env_idsdirectly.Suggested fix
- env_ids=[EnvironmentId(e) for e in self.env_ids], + env_ids=self.env_ids,
68-68: EagerSubtensorinitialization performs blocking network I/O in__init__.The previous lazy-loaded property deferred this cost. Now
Subtensor(network=...)blocks during construction, which could stall the event loop ifScheduler.__init__is ever called from an async context. Consider whether this should be deferred tostart()or wrapped inasyncio.to_thread. As per coding guidelines,kinitro/**/*.pyshould "Useasync deffor I/O and DB operations."kinitro/scoring/pareto.py (2)
7-8: Extraneous blank line after import.Line 8 is an empty changed line between the import and the existing blank line at line 9, resulting in a double blank line.
Suggested fix
from kinitro.types import EnvironmentId, MinerScores, MinerUID -
176-192: Consider O(1) UID lookup instead oflist.index().
uid_mapping.index(uid)is O(n). If these functions are called frequently (e.g., for every miner), building a reverse mapping{uid: idx}once would be more efficient. Not critical at current scale, but worth noting.Also applies to: 195-211
kinitro/scheduler/scoring.py (1)
188-205:"unknown"default should be wrapped asHotkey("unknown").On Line 190, the fallback
"unknown"is a plainstr, but it ends up as a value in aMinerScoreDatadict where thehotkeyfield is typed asHotkey. Strict type checkers will flag this.Proposed fix
- hotkey = miners_by_uid.get(uid, "unknown") + hotkey = miners_by_uid.get(uid, Hotkey("unknown"))kinitro/types.py (2)
16-25:NewTypealiases and composite type aliases are clear and useful.
MinerScores,MinerThresholds, andMinerFirstBlocksas plain type aliases are convenient. For stricter type-checker support, you could annotate them withTypeAlias(e.g.,MinerScores: TypeAlias = dict[...]), but this is not required.
155-161: ConsiderAnyinstead ofobjectfordomain_randomizationvalues.
dict[str, object]prevents attribute access on values without casting, which may be cumbersome for consumers. If the contents are truly heterogeneous,Anyis the pragmatic choice (with a justifying comment per coding guidelines). If the intent is to disallow arbitrary attribute access,objectis fine.kinitro/environments/genesis/base.py (1)
751-757: Barelisttype annotation lacks element type information.The type
dict[str, EncodedImage | list]uses barelist, which is overly broad. TheObservationmodel defines the same type for backward compatibility with legacy nested list formats (seedecode_image()function). However, consider narrowing the local variable todict[str, EncodedImage | list[Any]]for clarity, or drop the| listunion entirely if this method never assigns list values—currently onlyEncodedImageobjects are assigned viaencode_image().
rayokamoto
force-pushed
the
feat/typing
branch
3 times, most recently
from
11617f5 to
1bd4324
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
environments/_template/env.py (1)
206-213:⚠️ Potential issue | 🟡 Minor
_build_error_resultstill usesenv_id: str— should beEnvironmentId.Every call site passes an
EnvironmentId, and the rest of this file has been migrated. This parameter was missed.Proposed fix
def _build_error_result( self, - env_id: str, + env_id: EnvironmentId, task_id: int, seed: int, start_time: float, error: str, ) -> dict:kinitro/environments/genesis/scene_generator.py (1)
143-143:⚠️ Potential issue | 🟠 MajorUse
.valueto preserve the object_id format.When
obj_typeis anObjectTypeenum,f"{obj_type}"produces"ObjectType.BOX"instead of"box". This changes the object_id from"obj_00_red_box"to"obj_00_red_ObjectType.BOX", breaking downstream lookups inobject_statesdict and reward/success computation functions.Proposed fix
objects.append( SceneObjectConfig( - object_id=f"obj_{i:02d}_{color_name}_{obj_type}", + object_id=f"obj_{i:02d}_{color_name}_{obj_type.value}", object_type=obj_type,kinitro/backend/storage.py (1)
264-282:⚠️ Potential issue | 🟡 MinorType annotation mismatch:
weightsparameter should be converted before storing in JSON column.
weights: dict[MinerUID, float]is assigned directly toweights_json=weights(line 276), but the ORM model definesweights_json: Mapped[dict[str, float]](models.py line 124). SinceMinerUIDkeys areintat runtime, they serialize to JSON strings{"1": 0.5}. This works at runtime and the API route handles it viaint(k)on deserialization (routes/weights.py line 27), but it creates a type annotation mismatch.Convert the dict before storing to match the declared type:
weights_json={str(k): v for k, v in weights.items()}
🤖 Fix all issues with AI agents
In `@kinitro/scheduler/scoring.py`:
- Around line 172-176: The parameter type for miners_by_uid in
convert_to_scores_data is incorrect: change its annotation from dict[MinerUID,
str] to dict[MinerUID, Hotkey] (and update any imports) so it aligns with how
the map is used and with MinerScoreData.hotkey; also ensure the
default/placeholder usage (e.g., Hotkey("unknown")) remains a Hotkey instance
and update any related type hints or docstrings; verify call sites like the one
in kinitro/scheduler/main.py (which passes m.hotkey) still match the updated
signature.
In `@kinitro/validator/main.py`:
- Around line 166-171: The Subtensor websocket opened by Subtensor(network=...)
must be explicitly closed to avoid leaks: in the try/finally where you set
validator.subtensor and await validator.run(), update the finally block to first
check if validator.subtensor is set (not None) and then await its async close
method (e.g., await validator.subtensor.close() or the appropriate async cleanup
call provided by the Bittensor SDK), handling any exceptions, before calling
await validator.close(); reference the Validator instance and its
validator.subtensor attribute to locate where to add this explicit shutdown.
🧹 Nitpick comments (7)
environments/_template/env.py (1)
40-44: Missing return type annotation on_get_env.The other implementations (e.g.,
environments/genesis/env.py:66) annotate the return as-> RoboticsEnvironment. Consider adding the return type here for consistency with the template's role as a reference implementation. As per coding guidelines, "Add return types to public functions and methods."kinitro/executor/worker.py (1)
59-61: Consider extracting the family-from-env_id logic into a shared utility.The
env_id.split("/", maxsplit=1)[0] if "/" in env_id else env_idpattern is now duplicated inWorker._get_family(here) andExecutorConfig.get_image_for_env(config.pyline 218). A small helper (e.g., inkinitro/types.py) would reduce duplication.environments/metaworld/env.py (2)
317-319:_build_error_resultusesenv_id: strinstead ofEnvironmentId.All other methods in this class (
_get_env,evaluate,_run_evaluation) now useEnvironmentId, but_build_error_resultstill declaresenv_id: str. This should be updated for consistency with the rest of the typing migration.Proposed fix
def _build_error_result( self, - env_id: str, + env_id: EnvironmentId, task_id: int,
62-66: Missing return type annotation on_get_env.As per the updated AGENTS.md guideline: "Add type hints to all function signatures (parameters and return types), not just public ones." The return type should be annotated here. The same applies to
cleanupon line 343.environments/genesis/env.py (2)
324-327:_build_error_resultstill usesstrforenv_id— inconsistent with the rest of the file.Every other method in this file was updated to use
EnvironmentId, but_build_error_resultat Line 326 still declaresenv_id: str. Since all callers pass anEnvironmentId, this should match.Suggested fix
def _build_error_result( self, - env_id: str, + env_id: EnvironmentId, task_id: int,
62-62:_env_locksdict key type should beEnvironmentIdfor consistency.The keys stored in
_env_locksareEnvironmentIdvalues (see Line 162:self._env_locks.setdefault(env_id, ...)), but the annotation still saysdict[str, asyncio.Lock].Suggested fix
- self._env_locks: dict[str, asyncio.Lock] = {} + self._env_locks: dict[EnvironmentId, asyncio.Lock] = {}kinitro/types.py (1)
146-163: Consider whethernp.ndarrayannotations are sufficient here.
RobotStateDictandProceduralTaskResultuse barenp.ndarraywithout shape or dtype hints. This is standard practice (numpy's type stubs don't carry that info ergonomically yet), but worth noting that consumers must rely on conventions for expected shapes (e.g.,base_posis 3D,base_quatis 4D). A brief inline comment on expected shapes could help future readers.
rayokamoto
force-pushed
the
feat/typing
branch
from
1bd4324 to
7ab9bc4
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
pyproject.toml (1)
59-65:⚠️ Potential issue | 🟡 MinorDuplicate
devdependency groups with inconsistent version constraints.
[project.optional-dependencies] dev(lines 59-65) and[dependency-groups] dev(lines 116-121) declare overlapping packages with diverging lower bounds:
Package optional-dep dependency-group pytest >=9.0.0 >=9.0.2 pytest-asyncio >=0.21.0 >=1.3.0 ruff >=0.1.0 >=0.14.13 This can cause different resolved environments depending on which mechanism a developer uses (
pip install -e ".[dev]"vs. PEP 735dependency-groups). Consolidate into one source of truth or remove the redundant section.Also applies to: 116-121
kinitro/executor/verification.py (1)
584-603:⚠️ Potential issue | 🟡 MinorPotential type mismatch:
infofromenv.step()is MetaWorld's dict, notStepInfo.At Line 590,
env.step(mw_action)returns MetaWorld's native info dict (which may contain keys like"success","near_object", etc. that don't matchStepInfo). Theinfovariable is initialized asStepInfo = {}on Line 584 but then overwritten with MetaWorld's return value on Line 590. This is then returned directly on Line 603.Since
StepInfois aTypedDict(total=False), assigning an arbitrary dict from MetaWorld to it won't cause a runtime error, but it's a type-safety gap — the returned dict may contain keys outside theStepInfoschema.kinitro/environments/metaworld_env.py (1)
584-603:⚠️ Potential issue | 🟡 Minor
infois annotated asStepInfobut overwritten with MetaWorld's native info dict.On Line 584,
infois typed asStepInfo = {}, but on Line 590 it's reassigned fromenv.step()which returns MetaWorld's own info dict containing keys outside theStepInfoschema (e.g.,near_object,grasp_success, etc.). The returned value on Line 603 is this MetaWorld dict, not a conformingStepInfo.Consider either:
- Constructing a proper
StepInfofrom MetaWorld's info before returning, or- Keeping the annotation as
dict[str, Any]for the internal variable and converting at the return site.Suggested approach: construct StepInfo at return
- done: bool = terminated or truncated - - if full_obs is None: - raise RuntimeError("MetaWorld step returned no observation") - return self._build_observation(full_obs), float(total_reward), done, info + done: bool = terminated or truncated + + if full_obs is None: + raise RuntimeError("MetaWorld step returned no observation") + + step_info: StepInfo = { + "success": self._episode_success, + } + return self._build_observation(full_obs), float(total_reward), done, step_infokinitro/scheduler/main.py (1)
351-367:⚠️ Potential issue | 🟠 MajorSubtensor connection is never closed — websocket leak.
run_schedulerinitializesSubtensorviascheduler.start()(line 115), but thefinallyblock only callsscheduler.stop()(sets a flag) andstorage.close(). TheSubtensorwebsocket connection is never cleaned up. The validator (run_validator) already handles this correctly with an explicitsubtensor.close()in itsfinallyblock.Suggested fix
try: await scheduler.start() finally: + if scheduler.subtensor is not None: + try: + await asyncio.to_thread(scheduler.subtensor.close) + except Exception: + logger.warning("subtensor_close_failed", exc_info=True) await scheduler.stop() await storage.close()
🤖 Fix all issues with AI agents
In `@kinitro/scheduler/main.py`:
- Line 275: The dict comprehension is needlessly unwrapping and rewrapping
MinerUIDs; since compute_weights returns dict[ MinerUID, float ], change the
comprehension that builds weights (currently using MinerUID(int(k))) to simply
preserve the existing MinerUID keys and cast values to float (i.e., use the
existing weights variable from compute_weights and map v -> float(v)); update
the line building weights in the function where weights is used so keys remain
MinerUID and values are floats.
🧹 Nitpick comments (6)
tests/unit/test_crypto.py (1)
63-116: Consider annotating thekeypairfixture parameter.Since this PR is about improving typing, the
keypairparameter on these test methods (e.g., Line 63, 68, 76, 84, 93, 102, 110) could also be annotated askeypair: BackendKeypairfor consistency. The same applies tokeypairparams inTestEncryptDecryptandTestIntegration.kinitro/scoring/threshold.py (1)
43-43: Consider updatingepisodes_per_envdict type to useEnvironmentId.The
dict[str, int]branch ofepisodes_per_envcould bedict[EnvironmentId, int]for consistency with the rest of the typing migration. Minor nit sinceNewTypeis transparent at runtime.Suggested change
- episodes_per_env: int | dict[str, int], + episodes_per_env: int | dict[EnvironmentId, int],This also requires adding
EnvironmentIdto the import on line 5.kinitro/scheduler/main.py (1)
162-164: PreferRuntimeErroroverassertfor runtime initialization guard.
assertstatements are stripped when Python runs with-O(optimized mode). The validator (kinitro/validator/main.pylines 67–68) usesraise RuntimeError(...)for the identical check. Using the same pattern here would be more robust and consistent.Suggested fix
async def _run_evaluation_cycle(self) -> None: """Run a single evaluation cycle.""" - assert self.subtensor is not None, "subtensor not initialized; call start() first" + if self.subtensor is None: + raise RuntimeError("subtensor not initialized; call start() first") start_time = time.time()kinitro/api/routes/miners.py (1)
29-29: Minor:miners_dictkey type could beMinerUIDfor consistency.The dict key uses raw
s.uid(int from ORM) while the rest of the file wraps UIDs withMinerUID(...). Considerdict[MinerUID, MinerInfo]and keying byMinerUID(s.uid).kinitro/scheduler/scoring.py (1)
167-167: Minor:list(uids)preservesMinerUIDtype but dict is annotatedlist[int].
weights_to_u16returnslist[MinerUID]for UIDs.list(uids)is stilllist[MinerUID], but the annotationdict[str, list[int]]sayslist[int]. Since this dict appears to be a serialization payload for the chain, consider either annotating asdict[str, list[int] | list[MinerUID]]or explicitly casting:"uids": [int(u) for u in uids].tests/unit/test_pareto.py (1)
123-139: Consider extracting repeated environment ID lists to module-level constants.Patterns like
[EnvironmentId("a"), EnvironmentId("b")]and(EnvironmentId("a"), EnvironmentId("b"))are repeated in almost every test. A module-level constant (e.g.,ENV_AB = [EnvironmentId("a"), EnvironmentId("b")]) would reduce noise and make the tests easier to scan.
rayokamoto
force-pushed
the
feat/typing
branch
3 times, most recently
from
a80aa7c to
fd76206
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@kinitro/chain/weights.py`:
- Around line 3-9: The third‑party imports are not sorted per Ruff I and the
local import is mixed in; reorder the imports so third‑party packages are
grouped and alphabetized (bittensor, bittensor_wallet, numpy, structlog) and
place the local import from kinitro.types after a blank line as the local group,
i.e. import Subtensor and Wallet first, then numpy as np, then structlog,
followed by the local EligibilityResult and MinerUID import.
In `@kinitro/cli/crypto_commands.py`:
- Around line 7-9: The third‑party imports in this module are not sorted per
Ruff I rules; reorder the third‑party block alphabetically (e.g., import
bittensor.Subtensor, bittensor_wallet.Wallet, typer or equivalent alphabetical
order for the actual import names) and ensure import groups follow: standard
library first, then third‑party, then local, with all imports at the top of the
file; update the import statements referencing typer, Subtensor, and Wallet
accordingly so they follow Ruff I ordering.
In `@kinitro/cli/miner/commitment.py`:
- Around line 3-5: Reorder the third-party imports at the top of commitment.py
to satisfy Ruff I: group and sort them alphabetically so bittensor (Subtensor)
comes first, then bittensor_wallet (Wallet), then typer; ensure imports remain
at the module top (after any docstring) and preserve the referenced symbols
Subtensor, Wallet, and typer.
🧹 Nitpick comments (3)
kinitro/cli/testing_commands.py (1)
21-25: Consider adding return type annotation.Since this PR focuses on improving typing, consider adding
-> Nonetotest_scoring(andmock_minerbelow) to complete the typing coverage for public CLI commands.def test_scoring( n_miners: int = typer.Option(5, help="Number of simulated miners"), n_envs: int = typer.Option(3, help="Number of environments"), episodes_per_env: int = typer.Option(50, help="Simulated episodes per environment"), -): +) -> None:As per coding guidelines: "Add return types to public functions and methods."
kinitro/api/routes/miners.py (1)
29-46: Consider caching the MinerUID conversion to avoid repeated wrapping.Line 44 creates a new
MinerUID(s.uid)for each score when calculating averages. While functionally correct, this could be optimized by reusing the already-converted UIDs from the aggregation loop.Suggested optimization
# Calculate average success rate per miner for uid, miner in miners_dict.items(): - miner_scores = [s.success_rate for s in scores if MinerUID(s.uid) == uid] + miner_scores = [s.success_rate for s in scores if s.uid == int(uid)] if miner_scores: miner.avg_success_rate = sum(miner_scores) / len(miner_scores)kinitro/types.py (1)
165-172: Consider adding a comment explaining theAnyusage.Per coding guidelines,
Anyshould be avoided unless required, with an explanation in code comments. Heredomain_randomizationusesAnybecause the randomization parameters are open-ended and environment-specific.Suggested documentation
class ProceduralTaskResult(TypedDict): """Result of ProceduralTaskGenerator.generate().""" object_positions: np.ndarray # (n_objects, 3) target_positions: np.ndarray # (n_objects, 3) physics_params: dict[str, float] - domain_randomization: dict[str, Any] + domain_randomization: dict[str, Any] # Any: env-specific randomization params vary by implementation
rayokamoto
force-pushed
the
feat/typing
branch
from
fd76206 to
2deecb6
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
kinitro/rl_interface.py (1)
106-115:⚠️ Potential issue | 🟡 MinorAvoid
Any/rawlistin legacy image typings.
listandlist[Any]bypass the repo’s typing rules and make downstream validation harder. Consider a concrete legacy image alias (or add a brief justification ifAnyis truly required).As per coding guidelines, avoid Any type annotations unless required and explain why in code comments.♻️ Suggested typing tweak
+LegacyImage = list[list[int]] | list[list[list[int]]] + -def decode_image(encoded: EncodedImage | list) -> np.ndarray: +def decode_image(encoded: EncodedImage | LegacyImage) -> np.ndarray: @@ - rgb: dict[str, EncodedImage | list[Any]] = Field(default_factory=dict) - depth: dict[str, EncodedImage | list[Any]] = Field(default_factory=dict) + rgb: dict[str, EncodedImage | LegacyImage] = Field(default_factory=dict) + depth: dict[str, EncodedImage | LegacyImage] = Field(default_factory=dict)Also applies to: 154-155
🤖 Fix all issues with AI agents
In `@kinitro/cli/miner/commitment.py`:
- Around line 106-107: The synchronous Subtensor construction
(Subtensor(network=network)) in the commit and show_commitment flows must be
replaced with async usage: refactor the logic inside the functions that perform
I/O (the commit flow and show_commitment) into async helpers (e.g., async def
_commit_async(...), async def _show_commitment_async(...)) that create
AsyncSubtensor (or use SubtensorApi(async_subtensor=True)) as an async context
manager, await the necessary calls, and return results; then invoke these
helpers from the CLI entry points using asyncio.run(...). Update references to
Subtensor in the commit routine and show_commitment to use AsyncSubtensor
context manager and ensure Wallet creation remains compatible with async calls
if needed.
In `@kinitro/cli/miner/deploy.py`:
- Around line 5-9: Reorder the third-party imports to follow Ruff I ordering and
grouping: keep all third-party imports together and sort them alphabetically
(basilica, bittensor, bittensor_wallet, huggingface_hub, typer) so that typer
follows the other third‑party modules; ensure imports remain at the top of the
file and no extra grouping or blank lines split these third‑party imports.
- Around line 466-468: The code uses synchronous Subtensor/Wallet in
commit_model(); refactor commit_model() into an async def that uses the async
on-chain APIs (either AsyncSubtensor or SubtensorApi(async_subtensor=True))
instead of Subtensor, create the Wallet/Hotkey using the async-compatible API,
and perform the commit calls via the async client; then invoke this from the CLI
by either wrapping the call in asyncio.run(commit_model(...)) inside the
existing miner_deploy() or convert miner_deploy() to async and use typer's async
support to await commit_model(). Update references to the old synchronous
symbols (Subtensor, commit_model, miner_deploy) accordingly so all network I/O
uses the async APIs.
In `@kinitro/types.py`:
- Around line 175-186: Change the AffinetesEnv Protocol so list_environments()
returns list[EnvironmentId] instead of Any and add an explanatory comment on
evaluate()’s use of Any to document that payloads are environment-specific;
update the method signature for list_environments in the AffinetesEnv class and
add a short comment above async def evaluate(...) explaining why its
kwargs/return use Any are intentional and cannot be tightened.
In `@pyproject.toml`:
- Around line 60-63: The new dev dependency minimums ("pytest>=9.0.2",
"pytest-asyncio>=1.3.0", "pytest-cov>=4.0.0", "ruff>=0.14.13") may be
incompatible with our CI and test/coverage configs; verify CI Python runtime is
>=3.10 (or lower the pytest-asyncio pin), confirm pytest 9's path-argument
behavior won't break test invocations, ensure coverage measurement still works
given pytest-cov 4 removal of multiprocessing support (or revert to a compatible
pytest-cov/coverage combo), and check ruff 0.14.13/isort ordering changes
against our import-order expectations — then update the version strings in
pyproject.toml (or adjust CI/test configs) accordingly before merging.
🧹 Nitpick comments (2)
tests/unit/test_cycle_isolation.py (1)
54-54: Consider adding parameter type annotation for consistency.Since the fixture returns
AsyncMock, you could add the type annotation to the parameter as well for full consistency:async def test_no_incomplete_cycles(self, mock_session: AsyncMock) -> None:Same applies to the other test methods using this fixture.
kinitro/scoring/winners_take_all.py (1)
180-186: Add a wildcard case for exhaustiveness.Since
subset_weight_schemeis now aSubsetWeightSchemeenum (no longer a string), the match statement should include acase _:branch to guard against future enum variants being added. Without it,subset_weightwould be unbound if a new variant is introduced.Proposed fix
match subset_weight_scheme: case SubsetWeightScheme.LINEAR: subset_weight = float(subset_size) case SubsetWeightScheme.EXPONENTIAL: subset_weight = float(2 ** (subset_size - 1)) case SubsetWeightScheme.EQUAL: subset_weight = 1.0 + case _: + raise ValueError(f"Unknown subset weight scheme: {subset_weight_scheme}")
rayokamoto
force-pushed
the
feat/typing
branch
2 times, most recently
from
3356570 to
cba093e
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unit/test_genesis.py (1)
57-76:⚠️ Potential issue | 🟡 MinorUse
dict[str, Any] | Noneto match TaskSpec's type signature.The bare
dictannotation implicitly meansdict[Any, Any]. However, sinceTaskSpec.initial_stateis explicitly typed asdict[str, Any], the helper should use the same type for consistency. UsingAnyhere is required becauseTaskSpecmandates it.🔧 Suggested update
def _make_task_spec( task_type: TaskType = TaskType.NAVIGATE, target_object_id: str = "obj_00_red_box", target_object_type: ObjectType = ObjectType.BOX, target_position: list[float] | None = None, destination_object_id: str | None = None, destination_position: list[float] | None = None, - initial_state: dict | None = None, + initial_state: dict[str, Any] | None = None, ) -> TaskSpec:
🤖 Fix all issues with AI agents
In `@kinitro/chain/commitments.py`:
- Around line 159-166: In _parse_commitment_result, replace the unexplained Any
with a narrower annotation (e.g., Union[dict, object, None] or a
TypedDict/Protocol representing the expected shapes) or, if Any is unavoidable,
add a concise comment above the function that explains the substrate result’s
possible shapes (e.g., a dict with keys like "commitment"/"block" or an object
exposing a .value attribute) and why a broad type is required; update the
function signature from _parse_commitment_result(result: Any) to the chosen
narrower type and ensure the comment mentions the exact fields/properties the
function reads so future reviewers understand the assumptions.
- Around line 13-25: Reorder the import block to follow Ruff I: group standard
library imports first (dataclasses, typing), then third‑party packages sorted
alphabetically (structlog, bittensor imports
AsyncSubtensor/NeuronInfo/Subtensor, bittensor_wallet.Wallet,
cryptography.x25519), and finally local kinitro imports (kinitro.crypto
decrypt_deployment_id/encrypt_deployment_id and kinitro.types
BlockNumber/Hotkey/MinerUID/ParsedCommitment); keep existing from ... import
forms and preserve names like decrypt_deployment_id, AsyncSubtensor, Wallet,
x25519 and ParsedCommitment so the rest of the module (e.g., any references to
these symbols) continues to work.
In `@kinitro/cli/miner/commitment.py`:
- Around line 52-62: The function _get_neurons_hotkey_async currently allows a
negative uid which will index from the end of neurons; add a guard that returns
None when uid is out of bounds by checking if uid is less than 0 or uid >=
len(neurons) before accessing neurons[uid]. Update the conditional that
currently only checks uid >= len(neurons) to instead check (uid < 0 or uid >=
len(neurons)) and keep the rest of the logic (using AsyncSubtensor, neurons, and
returning neurons[uid].hotkey) unchanged.
🧹 Nitpick comments (3)
kinitro/cli/service_commands.py (1)
16-28: Consider adding return type annotations to CLI command functions.Since this PR focuses on typing improvements, the CLI command functions (
api,scheduler,executor) could benefit from explicit-> Nonereturn type annotations to align with the coding guidelines. As per coding guidelines, "Add return types to public functions and methods."💡 Example for the `api` function
def api( host: str = typer.Option("0.0.0.0", help="Host to bind to"), port: int = typer.Option(8000, help="Port to bind to"), database_url: str = typer.Option( "postgresql://postgres:postgres@localhost:5432/kinitro", help="PostgreSQL connection URL", ), no_auth: bool = typer.Option( False, "--no-auth", help="Disable API key authentication for task endpoints.", ), log_level: str = typer.Option("INFO", help="Logging level"), -): +) -> None:Same pattern would apply to
schedulerandexecutorfunctions.kinitro/api/routes/miners.py (1)
31-40: Minor redundancy:MinerUID(s.uid)used twice.Line 31 creates
uid = MinerUID(s.uid), but line 34 wraps it again withuid=MinerUID(s.uid). Consider reusing the existinguidvariable for consistency.♻️ Suggested simplification
for s in scores: uid = MinerUID(s.uid) if uid not in miners_dict: miners_dict[uid] = MinerInfo( - uid=MinerUID(s.uid), + uid=uid, hotkey=Hotkey(s.hotkey), last_evaluated_block=cycle.block_number, avg_success_rate=0.0, environments_evaluated=[], )kinitro/chain/commitments.py (1)
258-263: Consider an async variant for chain I/O.
read_miner_commitmentsperforms network calls but is synchronous; per project guidance, preferasync defwithAsyncSubtensorto avoid blocking event loops.As per coding guidelines: “
kinitro/**/*.py: Useasync deffor I/O and DB operations. Useasyncio.run(...)at CLI boundaries. Useasynccontextmanagerand session managers for DB work.”
rayokamoto
force-pushed
the
feat/typing
branch
from
cba093e to
5ecc68b
Compare
2 participants
Summary by CodeRabbit
Type System Enhancements
API / Backward Compatibility
Environments & Robotics
Documentation
CLI / Tests