Fix single-machine cross-process registrations by alberthli · Pull Request #69 · rai-opensource/judo

alberthli · 2025-07-11T19:57:04Z

This PR addresses broken CLI-based custom registrations on a single machine for the following. This addresses #68.

Tasks
Task-specific overrides
Update CHANGELOG.md
Update docs to explain sharp edge that you shouldn't use relative imports. Also indicate that if you are using multiple machines, you should currently only use the hydra-based override system.

Previously, these overrides were only loosely tested in a single process, but the propagation to other processes was not investigated (we only investigated the hydra-based solution deeply). This PR does the following to resolve this issue:

It saves the custom task and config definitions to file
This file is opened by all other processes, retrieving the custom definition

This PR currently does NOT solve the following, but we will want to soon, especially for mobile robots:

custom task registrations across machines

This simple test case passes for me currently:

from dataclasses import dataclass

from judo.config import set_config_overrides
from judo.optimizers.cem import CrossEntropyMethodConfig
from judo.tasks import CylinderPush, CylinderPushConfig


@dataclass
class MyCylinderPushConfig(CylinderPushConfig):
    """Custom configuration for MyCylinderPush task."""

class MyCylinderPush(CylinderPush):
    """Custom task to test registration."""

set_config_overrides(
    "my_cylinder_push",
    CrossEntropyMethodConfig,
    {
        "num_nodes": 5,
    },
)

if __name__ == "__main__":
    from judo.cli import app
    from judo.tasks import get_registered_tasks, register_task

    # on running this app, you should see my_cylinder_push in the dropdown
    # you should also see that num_nodes is set to 5 for it, and all other
    # values are set to global defaults with no overrides (so the task will
    # not perform well, but it will run)
    register_task("my_cylinder_push", MyCylinderPush, MyCylinderPushConfig)
    tasks = get_registered_tasks()
    
    app()

To test on other custom code, do the following:

pip uninstall judo-rai  # uninstall any previous versions in your virtual env
pip install git+https://github.com/alberthli/judo.git@e2ebc4c

…gistration

…o fix-task-registration

pculbertson

This one needs to cook a bit longer, per slack discussion

pculbertson · 2025-07-28T21:55:24Z

+        return
+    try:
+        data = json.loads(_OVERRIDES_PATH.read_text())
+    except Exception:


style: per the Google guide, I don't like catch-all except blocks. I also try to print or log the exception somewhere if I'm suppressing it.

https://google.github.io/styleguide/pyguide.html#244-decision

pculbertson · 2025-07-28T21:57:30Z

+
+
+def get_run_id(env_var: str = "JUDO_TASK_RUN_ID") -> str:
+    """Return a per-process run-ID and whether it was freshly created."""


nit: this docstring seems stale

pculbertson · 2025-07-28T22:03:41Z

+        _registered_tasks[name] = (task_cls, cfg_cls)
+
+
+def _persist_ephemeral_registry() -> None:


nit: persist doesn't really capture what this function is doing. maybe sync or update?

pculbertson · 2025-07-28T22:07:14Z

 def register_task(name: str, task_type: Type[Task], task_config_type: Type[TaskConfig]) -> None:
-    """Registers a new task."""
+    """Register a new task with the Judo framework for this run only."""
    _registered_tasks[name] = (task_type, task_config_type)


Should there be a _load_ephemeral_registry call here? Otherwise registry entries only get read once from file whenever __init__ is imported (per thread?) right?

Alternatively, you could change the logic of _persist_ephemeral_registry() to first read in what's currently in the file, merge it with what's in _registered_tasks in this thread + write it back?

pculbertson · 2025-07-28T22:11:21Z

+    directory.mkdir(parents=True, exist_ok=True)
+    path = directory / f"{prefix}_{run_id}.json"
+
+    # remove stale siblings


Maybe you should print a warning that you're doing this? I can't imagine why someone would do this, but if they happened to write an unrelated file {prefix}_{something_else}.json it's going to get deleted here. That'd be a hell of a thing to debug lol

pculbertson · 2025-07-28T22:12:10Z

+        return
+    try:
+        data = json.loads(_CUSTOM_REGISTRY_PATH.read_text())
+    except Exception:


same nit re: bare exceptions. not sure this should be suppressed if it errors...

pculbertson · 2025-07-28T22:17:35Z

+
+    for name, info in data.items():
+        # load Task class
+        task_mod = info.get("task_mod")


This will always be None, right? I don't see task_mod anywhere else in the codebase - is it vestigial?

fix CLI-based task registration across processes

08a697d

alberthli requested review from pculbertson and slecleach as code owners July 11, 2025 19:57

alberthli added 11 commits July 11, 2025 13:25

also fix persistence across processes for overrides

2bdd1d3

attempt fix for registration involving relative imports

6b8d480

Merge branch 'main' into fix-task-registration

cdd09a0

Merge branch 'main' into fix-task-registration

bf8731f

Merge branch 'main' into fix-task-registration

95afacb

Merge branch 'main' of github.com:bdaiinstitute/judo into fix-task-re…

792da12

…gistration

update changelog

fbc79d1

Merge branch 'fix-task-registration' of github.com:alberthli/judo int…

14ce53c

…o fix-task-registration

update docs with sharp edges

399d736

clean up some boilerplate

e2ebc4c

minor changes to typing

405b37d

pculbertson requested changes Jul 28, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix single-machine cross-process registrations#69

Fix single-machine cross-process registrations#69
alberthli wants to merge 12 commits intorai-opensource:mainfrom
alberthli:fix-task-registration

alberthli commented Jul 11, 2025 •

edited

Loading

Uh oh!

pculbertson left a comment

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

pculbertson Jul 28, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def get_run_id(env_var: str = "JUDO_TASK_RUN_ID") -> str:
		"""Return a per-process run-ID and whether it was freshly created."""

		_registered_tasks[name] = (task_cls, cfg_cls)


		def _persist_ephemeral_registry() -> None:

Conversation

alberthli commented Jul 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pculbertson left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

alberthli commented Jul 11, 2025 •

edited

Loading