Update app.py

app.py

@@ -458,45 +458,37 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- Runtime sitecustomize wrapper (SAFE) ------------------------------------
+# ---- OPTIONAL: legacy sitecustomize fallback (kept for belt & suspenders) ----
 def _install_workspace_env_fallback(cwd_for_train: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
     Creates a sitecustomize.py in the training cwd that monkeypatches
-    rtdetrv2_pytorch.src.core.workspace.create to ensure cfg['_pymodule'] exists
-    and is a *module object* before delegating to the original create().
+    rtdetrv2_pytorch.src.core.workspace.create to gracefully handle missing
+    cfg['_pymodule'] by falling back to $RTDETR_PYMODULE or module_default.
     """
     sc_path = os.path.join(cwd_for_train, "sitecustomize.py")
     code = textwrap.dedent(f"""
-        import os, importlib, types
+        import os, importlib
         try:
+            mod_path = os.environ.get("RTDETR_PYMODULE", "{module_default}")
             ws_mod = importlib.import_module("rtdetrv2_pytorch.src.core.workspace")
             _orig_create = ws_mod.create
-
-            def _ensure_module_object(name_or_none: str) -> types.ModuleType:
-                target = name_or_none.strip() if isinstance(name_or_none, str) and name_or_none.strip() else os.environ.get("RTDETR_PYMODULE", "{module_default}")
-                try:
-                    return importlib.import_module(target)
-                except Exception:
-                    return importlib.import_module("{module_default}")
-
             def _safe_create(name, cfg, *args, **kwargs):
                 pm = None
                 try:
                     pm = cfg.get("_pymodule", None)
                 except Exception:
                     pm = None
-
-                if isinstance(pm, types.ModuleType):
+                if not pm:
+                    pm = os.environ.get("RTDETR_PYMODULE", "{module_default}")
+                try:
+                    importlib.import_module(pm)
+                except Exception:
+                    pass
+                try:
+                    cfg["_pymodule"] = pm
+                except Exception:
                     pass
-                else:
-                    mod = _ensure_module_object(pm)
-                    try:
-                        cfg["_pymodule"] = mod
-                    except Exception:
-                        pass
-
                 return _orig_create(name, cfg, *args, **kwargs)
-
             ws_mod.create = _safe_create
         except Exception:
             pass
@@ -505,14 +497,99 @@ def _install_workspace_env_fallback(cwd_for_train: str, module_default: str = "r
         f.write(code)
     return sc_path
 
+# ---- NEW: direct, deterministic patch of workspace.create (V2) ---------------
+def _patch_workspace_create(repo_root: str, module_default: str = "rtdetrv2_pytorch.src") -> str | None:
+    """
+    Idempotently patches third_party/RT-DETRv2/rtdetrv2_pytorch/src/core/workspace.py
+    so that workspace.create() guarantees cfg['_pymodule'] is a *module object*,
+    not a string. If missing or string, it imports the module and writes it back.
+    Returns the path of the file patched (or None).
+    """
+    ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
+    if not os.path.isfile(ws_path):
+        return None
+
+    # If previously patched with V2, nothing to do
+    with open(ws_path, "r", encoding="utf-8") as f:
+        current_src = f.read()
+    bak_path = ws_path + ".bak"
+    if "_SAFE_CREATE_PATCH_V2" in current_src:
+        return ws_path
+
+    # If there was an older patch, restore from .bak to re-apply cleanly
+    if "_SAFE_CREATE_PATCH" in current_src and os.path.exists(bak_path):
+        try:
+            shutil.copy2(bak_path, ws_path)
+        except Exception:
+            pass
+        with open(ws_path, "r", encoding="utf-8") as f:
+            current_src = f.read()
+
+    # Backup original once
+    if not os.path.exists(bak_path):
+        try:
+            shutil.copy2(ws_path, bak_path)
+        except Exception:
+            pass
+
+    marker = "_SAFE_CREATE_PATCH_V2"
+    helper_code = f"""
+# --- {marker}: begin ---
+import importlib as _importlib
+import types as _types
+
+def _ensure_pymodule_object(_cfg, _default="{module_default}"):
+    pm = None
+    try:
+        pm = _cfg.get("_pymodule", None)
+    except Exception:
+        pm = None
+    if isinstance(pm, str) or pm is None:
+        name = pm if isinstance(pm, str) and pm.strip() else _default
+        try:
+            mod = _importlib.import_module(name)
+        except Exception:
+            mod = _importlib.import_module(_default)
+        try:
+            _cfg["_pymodule"] = mod
+        except Exception:
+            pass
+        return mod
+    if isinstance(pm, _types.ModuleType):
+        return pm
+    try:
+        mod = _importlib.import_module(_default)
+        _cfg["_pymodule"] = mod
+        return mod
+    except Exception:
+        return pm
+# --- {marker}: end ---
+"""
+
+    # Insert helper at top if missing
+    src = current_src
+    if marker not in src:
+        src = helper_code + src
+
+    # Inject normalization line right after def create(
+    inject_line = "    cfg['_pymodule'] = _ensure_pymodule_object(cfg)\n"
+    lines = src.splitlines(keepends=True)
+    for i, line in enumerate(lines):
+        if line.strip().startswith("def create(") and "_ensure_pymodule_object" not in "".join(lines[i:i+5]):
+            lines.insert(i + 1, inject_line)
+            break
+    new_src = "".join(lines)
+
+    with open(ws_path, "w", encoding="utf-8") as f:
+        f.write(new_src)
+    return ws_path
+
 def _unpatch_workspace_create(repo_root: str):
-    """If a previous run modified workspace.py, restore it from its .bak (best-effort)."""
     ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
     bak_path = ws_path + ".bak"
     if os.path.exists(bak_path):
         try:
             shutil.copy2(bak_path, ws_path)
-            logging.info("Restored original workspace.py from backup.")
         except Exception:
             pass
 
@@ -855,9 +932,12 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     if not base_cfg:
         raise gr.Error("Could not find a matching RT-DETRv2 config in the repo (S/M/M*/L/X).")
 
-    # Ensure we are NOT using any previous on-disk patch. Rely on runtime wrapper.
-    _unpatch_workspace_create(REPO_DIR)
-    logging.info("Using runtime sitecustomize monkey-patch for workspace.create; no on-disk edits.")
+    # Apply deterministic patch so workspace.create never fails on '_pymodule'
+    patched = _patch_workspace_create(REPO_DIR, module_default="rtdetrv2_pytorch.src")
+    if patched:
+        logging.info(f"Patched workspace.create at: {patched}")
+    else:
+        logging.warning("workspace.create patch not applied (file missing or atypical).")
 
     data_yaml = os.path.join(dataset_path, "data.yaml")
     with open(data_yaml, "r", encoding="utf-8") as f:
@@ -889,7 +969,7 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     q = Queue()
     def run_train():
         try:
-            # Runtime helper that ensures cfg['_pymodule'] is a module object.
+            # Optional: legacy sitecustomize helper (kept; patch above already handles it)
             train_cwd = os.path.dirname(train_script)
             _install_workspace_env_fallback(train_cwd)