diff --git a/CHANGELOG.md b/CHANGELOG.md index 267716ab..083edc88 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/docs/source/interface/config_registration.md b/docs/source/interface/config_registration.md index 454e9962..cd23185d 100644 --- a/docs/source/interface/config_registration.md +++ b/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: diff --git a/judo/config.py b/judo/config.py index 98ce9cb3..562608bd 100644 --- a/judo/config.py +++ b/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() diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py index 0a7537f7..629331a7 100644 --- a/judo/tasks/__init__.py +++ b/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,8 +14,9 @@ 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), @@ -18,16 +24,85 @@ "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__ = [ diff --git a/judo/utils/registration.py b/judo/utils/registration.py new file mode 100644 index 00000000..ca90a49c --- /dev/null +++ b/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