Fix single-machine cross-process registrations

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Fixed
 * Fixed bug after spec changes where unnamed geoms were causing `judo` to crash (@alberthli, @pculbertson, #72)
+* Fixed single-machine programmatic cross-process config registration (@alberthli, #69)
+    * NOTE: this does NOT fix programmatic registrations across machines. If you use a multi-machine setup, we recommend using hydra-based config management.
 
 ## Dev
 Bump version to 0.0.4 in pyproject.toml (@alberthli, #73)

docs/source/interface/config_registration.md

@@ -11,6 +11,12 @@ if __name__ == "__main__":
     app()
 ```
 
+> ⚠️ **Warning** ⚠️
+>
+> We highly recommend using proper Python package structure when defining your custom tasks/optimizers, or to at least stay away from relative imports in your code. This can give Judo problems when trying to locate your custom module for registration!
+>
+> Further, The above method for custom task/optimizer/config registration only works in the single-machine case due to the way we're handling multi-processing. For the multi-machine case, we recommend using our `hydra`-based config management system, which sends a copy of the configuration to each individual machine. This also allows you to keep using the `judo` CLI.
+
 If you instead want to use the `judo` CLI command to start the app, you can register the task and optimizer using a `hydra` config. We provide a convenient interface to do so. Consider this example:
 ```yaml
 defaults:

judo/config.py

@@ -1,13 +1,69 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
+import importlib
+import inspect
+import json
+import threading
 import warnings
 from dataclasses import MISSING, dataclass, fields, is_dataclass
+from importlib.util import module_from_spec, spec_from_file_location
 from typing import Any
 
 import numpy as np
 
+from judo.utils.registration import get_run_id, make_ephemeral_path
+
 _OVERRIDE_REGISTRY: dict[type, dict[str, Any]] = {}
 
+# like in judo/tasks/__init__.py, we use a run ID to create a temporary file for overrides
+_run_id = get_run_id()
+_OVERRIDES_PATH = make_ephemeral_path("judo_overrides", _run_id)
+_override_lock = threading.Lock()
+
+
+def _load_ephemeral_overrides() -> None:
+    """Load any overrides previously registered in this run."""
+    if not _OVERRIDES_PATH.is_file():
+        return
+    try:
+        data = json.loads(_OVERRIDES_PATH.read_text())
+    except Exception:
+        return
+
+    for entry in data:
+        cls_mod = entry.get("class_mod")
+        cls_qn = entry["class_qn"]
+
+        if cls_mod and cls_mod.startswith("judo."):
+            # safe: import the package module (avoids circular file re-exec)
+            mod = importlib.import_module(cls_mod)
+            cls = getattr(mod, cls_qn)
+        else:
+            # fallback for __main__ or truly external classes
+            spec = spec_from_file_location(f"_judo_override_{cls_qn}", entry["class_src"])
+            assert spec is not None, f"Could not load spec for {entry['class_src']}"
+            mod = module_from_spec(spec)
+            assert spec.loader is not None, f"Loader for {entry['class_src']} is None"
+            spec.loader.exec_module(mod)
+            cls = getattr(mod, cls_qn)
+
+        _OVERRIDE_REGISTRY[cls] = entry["overrides"]
+
+
+def _persist_ephemeral_overrides() -> None:
+    """Persist the current state of the override registry to disk."""
+    with _override_lock:
+        serial = []
+        for _cls, ov_map in _OVERRIDE_REGISTRY.items():
+            src = inspect.getsourcefile(_cls)  # source file of the class
+            qn = _cls.__qualname__  # qualified name of the class
+            mod = _cls.__module__  # module name of the class
+            if src:
+                serial.append({"class_src": src, "class_mod": mod, "class_qn": qn, "overrides": ov_map})
+
+
+_load_ephemeral_overrides()
+
 
 @dataclass
 class OverridableConfig:
@@ -94,3 +150,6 @@ def set_config_overrides(
                 UserWarning,
                 stacklevel=2,
             )
+
+    # persist the overrides to disk so distinct processes can access them
+    _persist_ephemeral_overrides()

judo/tasks/__init__.py

@@ -1,6 +1,11 @@
 # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
 
-from typing import Dict, Tuple, Type
+import importlib
+import inspect
+import json
+import threading
+from importlib.util import module_from_spec, spec_from_file_location
+from typing import Type
 
 from judo.tasks.base import Task, TaskConfig
 from judo.tasks.caltech_leap_cube import CaltechLeapCube, CaltechLeapCubeConfig
@@ -9,25 +14,95 @@
 from judo.tasks.fr3_pick import FR3Pick, FR3PickConfig
 from judo.tasks.leap_cube import LeapCube, LeapCubeConfig
 from judo.tasks.leap_cube_down import LeapCubeDown, LeapCubeDownConfig
+from judo.utils.registration import get_run_id, make_ephemeral_path
 
-_registered_tasks: Dict[str, Tuple[Type[Task], Type[TaskConfig]]] = {
+_registered_tasks: dict[str, tuple[Type[Task], Type[TaskConfig]]] = {
     "cylinder_push": (CylinderPush, CylinderPushConfig),
     "cartpole": (Cartpole, CartpoleConfig),
     "fr3_pick": (FR3Pick, FR3PickConfig),
     "leap_cube": (LeapCube, LeapCubeConfig),
     "leap_cube_down": (LeapCubeDown, LeapCubeDownConfig),
     "caltech_leap_cube": (CaltechLeapCube, CaltechLeapCubeConfig),
 }
+_builtin_names = set(_registered_tasks.keys())
 
+# set a run ID for this process that is used to persist programmatic registrations
+_run_id = get_run_id()
+_CUSTOM_REGISTRY_PATH = make_ephemeral_path("judo_tasks", _run_id)
+_registry_lock = threading.Lock()
 
-def get_registered_tasks() -> Dict[str, Tuple[Type[Task], Type[TaskConfig]]]:
-    """Returns a dictionary of registered tasks."""
+
+def _load_ephemeral_registry() -> None:
+    """On import, pull in any user-registered tasks from the temp file."""
+    if not _CUSTOM_REGISTRY_PATH.is_file():
+        return
+    try:
+        data = json.loads(_CUSTOM_REGISTRY_PATH.read_text())
+    except Exception:
+        return
+
+    for name, info in data.items():
+        # load Task class
+        task_mod = info.get("task_mod")
+        if task_mod and not task_mod.startswith("__main__"):
+            mod = importlib.import_module(task_mod)  # package import preserves relative imports in that module
+            task_cls = getattr(mod, info["task_qn"])
+        else:
+            # fallback: load by path
+            spec = spec_from_file_location(f"_judo_task_{name}", info["task_src"])
+            assert spec and spec.loader, f"Could not load task module {info['task_src']}"
+            mod = module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            task_cls = getattr(mod, info["task_qn"])
+
+        # load Config class
+        cfg_mod = info.get("cfg_mod")
+        if cfg_mod and not cfg_mod.startswith("__main__"):
+            mod2 = importlib.import_module(cfg_mod)
+            cfg_cls = getattr(mod2, info["cfg_qn"])
+        else:
+            spec2 = spec_from_file_location(f"_judo_cfg_{name}", info["cfg_src"])
+            assert spec2 and spec2.loader, f"Could not load config module {info['cfg_src']}"
+            mod2 = module_from_spec(spec2)
+            spec2.loader.exec_module(mod2)
+            cfg_cls = getattr(mod2, info["cfg_qn"])
+
+        _registered_tasks[name] = (task_cls, cfg_cls)
+
+
+def _persist_ephemeral_registry() -> None:
+    """After any register_task(), write the current state of the registry to disk."""
+    with _registry_lock:
+        out: dict[str, dict[str, str]] = {}
+        for name, (task_cls, cfg_cls) in _registered_tasks.items():
+            if name in _builtin_names:
+                continue
+            task_src = inspect.getsourcefile(task_cls)
+            cfg_src = inspect.getsourcefile(cfg_cls)
+            if task_src and cfg_src:
+                out[name] = {
+                    "task_src": task_src,
+                    "task_qn": task_cls.__qualname__,
+                    "cfg_src": cfg_src,
+                    "cfg_qn": cfg_cls.__qualname__,
+                }
+        _CUSTOM_REGISTRY_PATH.parent.mkdir(parents=True, exist_ok=True)
+        _CUSTOM_REGISTRY_PATH.write_text(json.dumps(out, indent=2))
+
+
+# load any ephemeral registrations on import
+_load_ephemeral_registry()
+
+
+def get_registered_tasks() -> dict[str, tuple[Type[Task], Type[TaskConfig]]]:
+    """Get the currently registered tasks."""
     return _registered_tasks
 
 
 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)
+    _persist_ephemeral_registry()
 
 
 __all__ = [

judo/utils/registration.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved.
+
+import atexit
+import os
+import uuid
+from pathlib import Path
+
+
+def get_run_id(env_var: str = "JUDO_TASK_RUN_ID") -> str:
+    """Return a per-process run-ID and whether it was freshly created."""
+    run_id = os.environ.get(env_var)
+    new_run = run_id is None
+    if new_run:
+        run_id = uuid.uuid4().hex
+        os.environ[env_var] = run_id
+    return run_id
+
+
+def make_ephemeral_path(prefix: str, run_id: str, directory: str | Path = "/tmp") -> Path:
+    """Create / clean a temp JSON path and register its teardown."""
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    path = directory / f"{prefix}_{run_id}.json"
+
+    # remove stale siblings
+    for old in directory.glob(f"{prefix}_*.json"):
+        if old.name != path.name:
+            try:
+                old.unlink()
+            except OSError:
+                pass
+
+    # ensure cleanup on interpreter exit
+    def _cleanup() -> None:
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+
+    atexit.register(_cleanup)
+    return path