From 08a697dc48875c747ba853ac41474a38d45adfd7 Mon Sep 17 00:00:00 2001 From: Albert Li Date: Fri, 11 Jul 2025 12:50:12 -0700 Subject: [PATCH 1/7] fix CLI-based task registration across processes --- judo/tasks/__init__.py | 87 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 85 insertions(+), 2 deletions(-) diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py index 0a7537f7..c5f18366 100644 --- a/judo/tasks/__init__.py +++ b/judo/tasks/__init__.py @@ -1,5 +1,13 @@ # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. +import atexit +import inspect +import json +import os +import threading +import uuid +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path from typing import Dict, Tuple, Type from judo.tasks.base import Task, TaskConfig @@ -18,16 +26,91 @@ "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 = os.environ.get("JUDO_TASK_RUN_ID") +if _run_id is None: + _run_id = uuid.uuid4().hex + os.environ["JUDO_TASK_RUN_ID"] = _run_id + _REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") + + def _cleanup_registry_file() -> None: + """Remove the ephemeral registry file on exit.""" + try: + _REGISTRY_PATH.unlink() + except FileNotFoundError: + pass + + atexit.register(_cleanup_registry_file) +else: + _REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") # subsequent imports will use the same file + +_lock = threading.Lock() + + +def _load_ephemeral_registry() -> None: + """On import, pull in any user-registered tasks from the temp file.""" + if not _REGISTRY_PATH.is_file(): + return + try: + data = json.loads(_REGISTRY_PATH.read_text()) + except Exception: + return + + for name, info in data.items(): + # load Task class + spec = spec_from_file_location(f"_judo_task_{name}", info["task_src"]) + assert spec is not None, f"Could not load task module {info['task_src']}" + mod = module_from_spec(spec) + assert spec.loader is not None, f"Could not load task module {info['task_src']}" + spec.loader.exec_module(mod) + task_cls = getattr(mod, info["task_qn"]) + + # load Config class + spec2 = spec_from_file_location(f"_judo_cfg_{name}", info["cfg_src"]) + assert spec2 is not None, f"Could not load config module {info['cfg_src']}" + mod2 = module_from_spec(spec2) + assert spec2.loader is not None, f"Could not load config module {info['cfg_src']}" + 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 _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__, + } + _REGISTRY_PATH.parent.mkdir(parents=True, exist_ok=True) + _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]]]: - """Returns a dictionary of registered tasks.""" + """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__ = [ From 2bdd1d34c6bdec12fa0f53d22f5e6be477de336a Mon Sep 17 00:00:00 2001 From: Albert Li Date: Fri, 11 Jul 2025 13:25:45 -0700 Subject: [PATCH 2/7] also fix persistence across processes for overrides --- judo/config.py | 86 ++++++++++++++++++++++++++++++++++++++++++ judo/tasks/__init__.py | 40 ++++++++++++-------- 2 files changed, 111 insertions(+), 15 deletions(-) diff --git a/judo/config.py b/judo/config.py index 98ce9cb3..56221aa8 100644 --- a/judo/config.py +++ b/judo/config.py @@ -1,13 +1,96 @@ # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. +import atexit +import importlib +import inspect +import json +import os +import threading +import uuid import warnings from dataclasses import MISSING, dataclass, fields, is_dataclass +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path from typing import Any import numpy as np _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 = os.environ.get("JUDO_TASK_RUN_ID") +_new_run = _run_id is None +if _new_run: + _run_id = uuid.uuid4().hex + os.environ["JUDO_TASK_RUN_ID"] = _run_id +_OVERRIDES_PATH = Path(f"/tmp/judo_overrides_{_run_id}.json") + +# remove any potential stale judo_overrides_*.json files +for old in _OVERRIDES_PATH.parent.glob("judo_overrides_*.json"): + if old.name != _OVERRIDES_PATH.name: + try: + old.unlink() + except OSError: + pass + + +# register a cleanup function to remove the file on exit +def _cleanup_registry_file() -> None: + """Remove the ephemeral registry file on exit of the main.""" + try: + _OVERRIDES_PATH.unlink() + except FileNotFoundError: + pass + + +atexit.register(_cleanup_registry_file) + +_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 +177,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 c5f18366..80c8415e 100644 --- a/judo/tasks/__init__.py +++ b/judo/tasks/__init__.py @@ -30,31 +30,41 @@ # set a run ID for this process that is used to persist programmatic registrations _run_id = os.environ.get("JUDO_TASK_RUN_ID") -if _run_id is None: +_new_run = _run_id is None +if _new_run: _run_id = uuid.uuid4().hex os.environ["JUDO_TASK_RUN_ID"] = _run_id - _REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") +_CUSTOM_REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") - def _cleanup_registry_file() -> None: - """Remove the ephemeral registry file on exit.""" +# remove any potential stale judo_tasks_*.json files +for old in _CUSTOM_REGISTRY_PATH.parent.glob("judo_tasks_*.json"): + if old.name != _CUSTOM_REGISTRY_PATH.name: try: - _REGISTRY_PATH.unlink() - except FileNotFoundError: + old.unlink() + except OSError: pass - atexit.register(_cleanup_registry_file) -else: - _REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") # subsequent imports will use the same file -_lock = threading.Lock() +# register a cleanup function to remove the file on exit +def _cleanup_registry_file() -> None: + """Remove the ephemeral registry file on exit of the main.""" + try: + _CUSTOM_REGISTRY_PATH.unlink() + except FileNotFoundError: + pass + + +atexit.register(_cleanup_registry_file) + +_registry_lock = threading.Lock() def _load_ephemeral_registry() -> None: """On import, pull in any user-registered tasks from the temp file.""" - if not _REGISTRY_PATH.is_file(): + if not _CUSTOM_REGISTRY_PATH.is_file(): return try: - data = json.loads(_REGISTRY_PATH.read_text()) + data = json.loads(_CUSTOM_REGISTRY_PATH.read_text()) except Exception: return @@ -80,7 +90,7 @@ def _load_ephemeral_registry() -> None: def _persist_ephemeral_registry() -> None: """After any register_task(), write the current state of the registry to disk.""" - with _lock: + with _registry_lock: out: Dict[str, Dict[str, str]] = {} for name, (task_cls, cfg_cls) in _registered_tasks.items(): if name in _builtin_names: @@ -94,8 +104,8 @@ def _persist_ephemeral_registry() -> None: "cfg_src": cfg_src, "cfg_qn": cfg_cls.__qualname__, } - _REGISTRY_PATH.parent.mkdir(parents=True, exist_ok=True) - _REGISTRY_PATH.write_text(json.dumps(out, indent=2)) + _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 From 6b8d480e18c1fa0e8c61048f7a7f152903a761c1 Mon Sep 17 00:00:00 2001 From: Albert Li Date: Fri, 11 Jul 2025 15:08:48 -0700 Subject: [PATCH 3/7] attempt fix for registration involving relative imports --- judo/tasks/__init__.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py index 80c8415e..c155b280 100644 --- a/judo/tasks/__init__.py +++ b/judo/tasks/__init__.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. import atexit +import importlib import inspect import json import os @@ -70,20 +71,29 @@ def _load_ephemeral_registry() -> None: for name, info in data.items(): # load Task class - spec = spec_from_file_location(f"_judo_task_{name}", info["task_src"]) - assert spec is not None, f"Could not load task module {info['task_src']}" - mod = module_from_spec(spec) - assert spec.loader is not None, f"Could not load task module {info['task_src']}" - spec.loader.exec_module(mod) - task_cls = getattr(mod, info["task_qn"]) + 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 - spec2 = spec_from_file_location(f"_judo_cfg_{name}", info["cfg_src"]) - assert spec2 is not None, f"Could not load config module {info['cfg_src']}" - mod2 = module_from_spec(spec2) - assert spec2.loader is not None, f"Could not load config module {info['cfg_src']}" - spec2.loader.exec_module(mod2) - cfg_cls = getattr(mod2, info["cfg_qn"]) + 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) From fbc79d1f9cad8fc85d5e3d3b5239e615e95a1098 Mon Sep 17 00:00:00 2001 From: Albert Li Date: Wed, 23 Jul 2025 16:24:09 -0700 Subject: [PATCH 4/7] update changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) 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) From 399d736ecd63f4d1c44d09d85152f69a5c1da0ef Mon Sep 17 00:00:00 2001 From: Albert Li Date: Wed, 23 Jul 2025 16:34:20 -0700 Subject: [PATCH 5/7] update docs with sharp edges --- docs/source/interface/config_registration.md | 6 ++++++ 1 file changed, 6 insertions(+) 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: From e2ebc4c53d75e8932873d91c6beb6b35642c43b7 Mon Sep 17 00:00:00 2001 From: Albert Li Date: Wed, 23 Jul 2025 16:46:26 -0700 Subject: [PATCH 6/7] clean up some boilerplate --- judo/config.py | 35 ++++---------------------------- judo/tasks/__init__.py | 34 +++---------------------------- judo/utils/registration.py | 41 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 48 insertions(+), 62 deletions(-) create mode 100644 judo/utils/registration.py diff --git a/judo/config.py b/judo/config.py index 56221aa8..562608bd 100644 --- a/judo/config.py +++ b/judo/config.py @@ -1,50 +1,23 @@ # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. -import atexit import importlib import inspect import json -import os import threading -import uuid import warnings from dataclasses import MISSING, dataclass, fields, is_dataclass from importlib.util import module_from_spec, spec_from_file_location -from pathlib import Path 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 = os.environ.get("JUDO_TASK_RUN_ID") -_new_run = _run_id is None -if _new_run: - _run_id = uuid.uuid4().hex - os.environ["JUDO_TASK_RUN_ID"] = _run_id -_OVERRIDES_PATH = Path(f"/tmp/judo_overrides_{_run_id}.json") - -# remove any potential stale judo_overrides_*.json files -for old in _OVERRIDES_PATH.parent.glob("judo_overrides_*.json"): - if old.name != _OVERRIDES_PATH.name: - try: - old.unlink() - except OSError: - pass - - -# register a cleanup function to remove the file on exit -def _cleanup_registry_file() -> None: - """Remove the ephemeral registry file on exit of the main.""" - try: - _OVERRIDES_PATH.unlink() - except FileNotFoundError: - pass - - -atexit.register(_cleanup_registry_file) - +_run_id = get_run_id() +_OVERRIDES_PATH = make_ephemeral_path("judo_overrides", _run_id) _override_lock = threading.Lock() diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py index c155b280..962f9d37 100644 --- a/judo/tasks/__init__.py +++ b/judo/tasks/__init__.py @@ -1,14 +1,10 @@ # Copyright (c) 2025 Robotics and AI Institute LLC. All rights reserved. -import atexit import importlib import inspect import json -import os import threading -import uuid from importlib.util import module_from_spec, spec_from_file_location -from pathlib import Path from typing import Dict, Tuple, Type from judo.tasks.base import Task, TaskConfig @@ -18,6 +14,7 @@ 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]]] = { "cylinder_push": (CylinderPush, CylinderPushConfig), @@ -30,33 +27,8 @@ _builtin_names = set(_registered_tasks.keys()) # set a run ID for this process that is used to persist programmatic registrations -_run_id = os.environ.get("JUDO_TASK_RUN_ID") -_new_run = _run_id is None -if _new_run: - _run_id = uuid.uuid4().hex - os.environ["JUDO_TASK_RUN_ID"] = _run_id -_CUSTOM_REGISTRY_PATH = Path(f"/tmp/judo_tasks_{_run_id}.json") - -# remove any potential stale judo_tasks_*.json files -for old in _CUSTOM_REGISTRY_PATH.parent.glob("judo_tasks_*.json"): - if old.name != _CUSTOM_REGISTRY_PATH.name: - try: - old.unlink() - except OSError: - pass - - -# register a cleanup function to remove the file on exit -def _cleanup_registry_file() -> None: - """Remove the ephemeral registry file on exit of the main.""" - try: - _CUSTOM_REGISTRY_PATH.unlink() - except FileNotFoundError: - pass - - -atexit.register(_cleanup_registry_file) - +_run_id = get_run_id() +_CUSTOM_REGISTRY_PATH = make_ephemeral_path("judo_tasks", _run_id) _registry_lock = threading.Lock() 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 From 405b37da9427c362049745a898cf3630f5c8acb1 Mon Sep 17 00:00:00 2001 From: Albert Li Date: Wed, 23 Jul 2025 16:51:51 -0700 Subject: [PATCH 7/7] minor changes to typing --- judo/tasks/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/judo/tasks/__init__.py b/judo/tasks/__init__.py index 962f9d37..629331a7 100644 --- a/judo/tasks/__init__.py +++ b/judo/tasks/__init__.py @@ -5,7 +5,7 @@ import json import threading from importlib.util import module_from_spec, spec_from_file_location -from typing import Dict, Tuple, Type +from typing import Type from judo.tasks.base import Task, TaskConfig from judo.tasks.caltech_leap_cube import CaltechLeapCube, CaltechLeapCubeConfig @@ -16,7 +16,7 @@ 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), @@ -73,7 +73,7 @@ def _load_ephemeral_registry() -> None: 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]] = {} + out: dict[str, dict[str, str]] = {} for name, (task_cls, cfg_cls) in _registered_tasks.items(): if name in _builtin_names: continue @@ -94,7 +94,7 @@ def _persist_ephemeral_registry() -> None: _load_ephemeral_registry() -def get_registered_tasks() -> Dict[str, Tuple[Type[Task], Type[TaskConfig]]]: +def get_registered_tasks() -> dict[str, tuple[Type[Task], Type[TaskConfig]]]: """Get the currently registered tasks.""" return _registered_tasks