Create Hub artifact collections lazily (#245)

* Create Hub artifact collections lazily

Co-authored-by: OpenAI Codex <codex@openai.com>

* Cache lazy Hub collection slugs in child hooks

Co-authored-by: OpenAI Codex <codex@openai.com>

---------

Co-authored-by: OpenAI Codex <codex@openai.com>

agent/core/agent_loop.py

@@ -27,7 +27,6 @@ from agent.core.cost_estimation import CostEstimate, estimate_tool_cost
 from agent.messaging.gateway import NotificationGateway
 from agent.core import telemetry
 from agent.core.doom_loop import check_for_doom_loop
-from agent.core.hub_artifacts import start_session_artifact_collection_task
 from agent.core.llm_params import _resolve_llm_params
 from agent.core.prompt_caching import with_prompt_caching
 from agent.core.session import DEFAULT_SESSION_LOG_DIR, Event, OpType, Session
@@ -2024,7 +2023,6 @@ async def submission_loop(
     )
     if session_holder is not None:
         session_holder[0] = session
-    start_session_artifact_collection_task(session, token=hf_token)
     logger.info("Agent loop started")
 
     # Retry any failed uploads from previous sessions (fire-and-forget).

agent/core/hub_artifacts.py

@@ -1,6 +1,5 @@
 """Best-effort Hub metadata for artifacts generated by ML Intern sessions."""
 
-import asyncio
 import base64
 import logging
 import re
@@ -11,7 +10,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any
 
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import hf_hub_download
 from huggingface_hub.repocard import metadata_load, metadata_save
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
 
@@ -29,7 +28,6 @@ _UUID_SESSION_ID_RE = re.compile(
 _KNOWN_ARTIFACTS_ATTR = "_ml_intern_known_hub_artifacts"
 _REGISTERED_ARTIFACTS_ATTR = "_ml_intern_registered_hub_artifacts"
 _COLLECTION_SLUG_ATTR = "_ml_intern_artifact_collection_slug"
-_COLLECTION_TASK_ATTR = "_ml_intern_artifact_collection_task"
 _SESSION_ARTIFACT_SET_FALLBACK: dict[tuple[int, str], set[str]] = {}
 _USAGE_HEADING_RE = re.compile(
     r"^#{2,6}\s+(usage|how to use|using this (model|dataset)|use this (model|dataset))\b",
@@ -307,70 +305,6 @@ def _ensure_collection_slug(
     return slug
 
 
-async def ensure_session_artifact_collection(
-    session: Any,
-    *,
-    token: str | bool | None = None,
-) -> str | None:
-    """Create/cache the per-session artifact collection without raising."""
-    if session is None or not getattr(session, "session_id", None):
-        return None
-    token_value = token if token is not None else getattr(session, "hf_token", None)
-    if not token_value:
-        return None
-
-    try:
-        api = HfApi(token=token_value)
-        return await asyncio.to_thread(
-            _ensure_collection_slug,
-            api,
-            session,
-            token=token_value,
-        )
-    except Exception as e:
-        logger.warning(
-            "ML Intern session collection creation failed for %s: %s",
-            _safe_session_id(session),
-            e,
-        )
-        return None
-
-
-def start_session_artifact_collection_task(
-    session: Any,
-    *,
-    token: str | bool | None = None,
-) -> asyncio.Task | None:
-    """Schedule best-effort collection creation for a newly started session."""
-    if session is None or not getattr(session, "session_id", None):
-        return None
-    if getattr(session, _COLLECTION_SLUG_ATTR, None):
-        return None
-
-    token_value = token if token is not None else getattr(session, "hf_token", None)
-    if not token_value:
-        return None
-
-    existing = getattr(session, _COLLECTION_TASK_ATTR, None)
-    if isinstance(existing, asyncio.Task) and not existing.done():
-        return existing
-
-    try:
-        loop = asyncio.get_running_loop()
-    except RuntimeError:
-        return None
-
-    async def _run() -> None:
-        await ensure_session_artifact_collection(session, token=token_value)
-
-    task = loop.create_task(_run())
-    try:
-        setattr(session, _COLLECTION_TASK_ATTR, task)
-    except Exception:
-        logger.debug("Could not attach ML Intern collection task to session")
-    return task
-
-
 def _add_to_collection(
     api: Any,
     session: Any,
@@ -378,10 +312,10 @@ def _add_to_collection(
     repo_type: str,
     *,
     token: str | bool | None = None,
-) -> None:
+) -> bool:
     slug = _ensure_collection_slug(api, session, token=token)
     if not slug:
-        return
+        return False
     api.add_collection_item(
         collection_slug=slug,
         item_id=repo_id,
@@ -393,6 +327,7 @@ def _add_to_collection(
         exists_ok=True,
         token=token,
     )
+    return True
 
 
 def register_hub_artifact(
@@ -436,8 +371,13 @@ def register_hub_artifact(
         logger.debug("ML Intern repo-card update failed for %s: %s", repo_id, e)
 
     try:
-        _add_to_collection(api, session, repo_id, repo_type, token=token_value)
-        collection_updated = True
+        collection_updated = _add_to_collection(
+            api,
+            session,
+            repo_id,
+            repo_type,
+            token=token_value,
+        )
     except Exception as e:
         logger.debug("ML Intern collection update failed for %s: %s", repo_id, e)
 
@@ -490,6 +430,13 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                 re.IGNORECASE | re.MULTILINE,
             )
             front_matter_re = re.compile(r"\\A---\\s*\\n.*?\\n---\\s*\\n?", re.DOTALL)
+            collection_cache_path = (
+                os.environ.get("ML_INTERN_ARTIFACT_COLLECTION_CACHE")
+                or str(
+                    Path(tempfile.gettempdir())
+                    / f"ml-intern-artifacts-{{session_id}}.collection"
+                )
+            )
 
             def _token(value=None, api=None):
                 if isinstance(value, str) and value:
@@ -602,6 +549,15 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                 nonlocal collection_slug
                 if collection_slug:
                     return collection_slug
+                try:
+                    cached_slug = Path(collection_cache_path).read_text(
+                        encoding="utf-8"
+                    ).strip()
+                    if cached_slug:
+                        collection_slug = cached_slug
+                        return collection_slug
+                except Exception:
+                    pass
                 collection = api.create_collection(
                     title=collection_title,
                     description=(
@@ -613,6 +569,13 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                     token=token_value,
                 )
                 collection_slug = getattr(collection, "slug", None)
+                if collection_slug:
+                    try:
+                        cache_path = Path(collection_cache_path)
+                        cache_path.parent.mkdir(parents=True, exist_ok=True)
+                        cache_path.write_text(collection_slug, encoding="utf-8")
+                    except Exception:
+                        pass
                 return collection_slug
 
             def _register(
@@ -637,6 +600,7 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                 try:
                     token_value = _token(token_value)
                     api = HfApi(token=token_value)
+                    card_updated = False
                     try:
                         current = _readme(api, repo_id, repo_type, token_value)
                         updated = _augment(
@@ -652,8 +616,10 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                                 token=token_value,
                                 commit_message="Update ML Intern artifact metadata",
                             )
+                        card_updated = True
                     except Exception:
                         pass
+                    collection_updated = False
                     try:
                         slug = _ensure_collection(api, token_value)
                         if slug:
@@ -668,9 +634,11 @@ def build_hub_artifact_sitecustomize(session: Any) -> str:
                                 exists_ok=True,
                                 token=token_value,
                             )
+                            collection_updated = True
                     except Exception:
                         pass
-                    registered.add(key)
+                    if card_updated and collection_updated:
+                        registered.add(key)
                 finally:
                     registering = False
 

backend/session_manager.py

@@ -12,7 +12,6 @@ from typing import Any, Optional
 
 from agent.config import load_config
 from agent.core.agent_loop import process_submission
-from agent.core.hub_artifacts import start_session_artifact_collection_task
 from agent.core.session import Event, OpType, Session
 from agent.core.session_persistence import get_session_store
 from agent.core.tools import ToolRouter
@@ -136,7 +135,6 @@ class SessionManager:
         self.sessions: dict[str, AgentSession] = {}
         self._lock = asyncio.Lock()
         self.persistence_store = None
-        self.enable_hub_artifact_collections = True
 
     async def start(self) -> None:
         """Start shared background resources."""
@@ -413,28 +411,6 @@ class SessionManager:
         session.sandbox_preload_cancel_event = None
         self._start_cpu_sandbox_preload(agent_session)
 
-    def _start_hub_artifact_collection(self, agent_session: AgentSession) -> None:
-        """Kick off best-effort Hub collection creation for the session."""
-        if not getattr(self, "enable_hub_artifact_collections", False):
-            return
-        session = agent_session.session
-        if not getattr(session, "session_id", None):
-            try:
-                session.session_id = agent_session.session_id
-            except Exception:
-                logger.debug("Could not attach session id for Hub artifact collection")
-        token = agent_session.hf_token or getattr(session, "hf_token", None)
-        if not token:
-            return
-        try:
-            start_session_artifact_collection_task(session, token=token)
-        except Exception as e:
-            logger.debug(
-                "Failed to schedule Hub artifact collection for %s: %s",
-                agent_session.session_id,
-                e,
-            )
-
     async def _clear_persisted_sandbox_metadata(self, session_id: str) -> None:
         try:
             await self._store().update_session_fields(
@@ -591,7 +567,6 @@ class SessionManager:
                     existing,
                     preload_sandbox=preload_sandbox,
                 )
-                self._start_hub_artifact_collection(existing)
                 return existing
             return None
 
@@ -613,7 +588,6 @@ class SessionManager:
                     existing,
                     preload_sandbox=preload_sandbox,
                 )
-                self._start_hub_artifact_collection(existing)
                 return existing
             return None
 
@@ -700,9 +674,7 @@ class SessionManager:
                 hf_token=hf_token,
                 hf_username=hf_username,
             )
-            self._start_hub_artifact_collection(started)
             return started
-        self._start_hub_artifact_collection(agent_session)
         if preload_sandbox:
             self._start_cpu_sandbox_preload(agent_session)
         logger.info("Restored session %s for user %s", session_id, owner or user_id)
@@ -785,7 +757,6 @@ class SessionManager:
             event_queue=event_queue,
             tool_router=tool_router,
         )
-        self._start_hub_artifact_collection(agent_session)
         await self.persist_session_snapshot(agent_session, runtime_state="idle")
         self._start_cpu_sandbox_preload(agent_session)
 

tests/unit/test_hub_artifacts.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 from types import SimpleNamespace
 
@@ -11,12 +10,10 @@ from agent.core.hub_artifacts import (
     artifact_collection_title,
     augment_repo_card_content,
     build_hub_artifact_sitecustomize,
-    ensure_session_artifact_collection,
     is_known_hub_artifact,
     is_sandbox_hub_repo,
     register_hub_artifact,
     remember_hub_artifact,
-    start_session_artifact_collection_task,
     wrap_shell_command_with_hub_artifact_bootstrap,
 )
 from agent.tools import local_tools, sandbox_tool
@@ -207,6 +204,7 @@ def test_register_hub_artifact_retries_after_partial_failure(monkeypatch):
     def add_to_collection(*args, **kwargs):
         nonlocal collection_attempts
         collection_attempts += 1
+        return True
 
     monkeypatch.setattr(
         hub_artifacts,
@@ -238,6 +236,7 @@ def test_register_hub_artifact_retries_after_collection_failure(monkeypatch):
         collection_attempts += 1
         if collection_attempts == 1:
             raise RuntimeError("temporary collection failure")
+        return True
 
     monkeypatch.setattr(hub_artifacts, "_update_repo_card", update_repo_card)
     monkeypatch.setattr(
@@ -271,63 +270,6 @@ def test_session_artifact_set_falls_back_when_session_rejects_attrs(caplog):
     assert "using process-local fallback state" in caplog.text
 
 
-@pytest.mark.asyncio
-async def test_ensure_session_artifact_collection_uses_user_token(monkeypatch):
-    session = _session()
-    calls = []
-
-    class FakeApi:
-        def __init__(self, token):
-            self.token = token
-
-    def fake_ensure_collection_slug(api, seen_session, **kwargs):
-        calls.append((api.token, seen_session, kwargs))
-        return "alice/ml-intern-artifacts"
-
-    monkeypatch.setattr(hub_artifacts, "HfApi", FakeApi)
-    monkeypatch.setattr(
-        hub_artifacts,
-        "_ensure_collection_slug",
-        fake_ensure_collection_slug,
-    )
-
-    slug = await ensure_session_artifact_collection(session, token="hf-token")
-
-    assert slug == "alice/ml-intern-artifacts"
-    assert calls == [
-        ("hf-token", session, {"token": "hf-token"}),
-    ]
-
-
-@pytest.mark.asyncio
-async def test_start_session_artifact_collection_task_dedupes(monkeypatch):
-    session = _session()
-    calls = []
-
-    async def fake_ensure_session_artifact_collection(seen_session, **kwargs):
-        calls.append((seen_session, kwargs))
-        await asyncio.sleep(0)
-        return "alice/ml-intern-artifacts"
-
-    monkeypatch.setattr(
-        hub_artifacts,
-        "ensure_session_artifact_collection",
-        fake_ensure_session_artifact_collection,
-    )
-
-    task = start_session_artifact_collection_task(session, token="hf-token")
-    second = start_session_artifact_collection_task(session, token="hf-token")
-
-    assert task is not None
-    assert second is task
-    await task
-    assert calls == [(session, {"token": "hf-token"})]
-
-
-def test_start_session_artifact_collection_task_skips_without_token():
-    assert start_session_artifact_collection_task(_session()) is None
-
-
 @pytest.mark.asyncio
 async def test_hf_repo_git_create_repo_registers_artifact(monkeypatch):
     session = _session()
@@ -535,6 +477,78 @@ def test_sitecustomize_bootstrap_reuses_existing_collection_slug():
     )
 
 
+def test_sitecustomize_caches_lazy_collection_slug_across_bootstraps(
+    monkeypatch,
+    tmp_path,
+):
+    import huggingface_hub as hub
+    from huggingface_hub import HfApi
+
+    readme_path = tmp_path / "README.md"
+    readme_path.write_text("# Existing Model\n", encoding="utf-8")
+    cache_path = tmp_path / "collection-slug.txt"
+    collection_slug = "alice/ml-intern-artifacts-2026-05-05-session-123"
+    uploads = []
+    downloads = []
+    collection_creates = []
+    collection_items = []
+
+    def fake_upload_file(self, **kwargs):
+        uploads.append(kwargs)
+        return SimpleNamespace()
+
+    def fake_hf_hub_download(*args, **kwargs):
+        downloads.append((args, kwargs))
+        return str(readme_path)
+
+    def fake_create_collection(self, **kwargs):
+        collection_creates.append(kwargs)
+        return SimpleNamespace(slug=collection_slug)
+
+    def fake_add_collection_item(self, **kwargs):
+        collection_items.append(kwargs)
+
+    monkeypatch.setenv("ML_INTERN_ARTIFACT_COLLECTION_CACHE", str(cache_path))
+    code = build_hub_artifact_sitecustomize(_session())
+
+    def install_fresh_bootstrap():
+        monkeypatch.setattr(HfApi, "upload_file", fake_upload_file)
+        monkeypatch.setattr(HfApi, "create_collection", fake_create_collection)
+        monkeypatch.setattr(HfApi, "add_collection_item", fake_add_collection_item)
+        monkeypatch.setattr(hub, "hf_hub_download", fake_hf_hub_download)
+        exec(code, {})
+        assert HfApi.upload_file is not fake_upload_file
+
+    install_fresh_bootstrap()
+    HfApi(token="hf-token").upload_file(
+        path_or_fileobj=b"weights",
+        path_in_repo="model.safetensors",
+        repo_id="alice/model-a",
+        repo_type="model",
+        token="hf-token",
+    )
+
+    install_fresh_bootstrap()
+    HfApi(token="hf-token").upload_file(
+        path_or_fileobj=b"weights",
+        path_in_repo="model.safetensors",
+        repo_id="alice/model-b",
+        repo_type="model",
+        token="hf-token",
+    )
+
+    assert cache_path.read_text(encoding="utf-8") == collection_slug
+    assert len(collection_creates) == 1
+    assert [item["item_id"] for item in collection_items] == [
+        "alice/model-a",
+        "alice/model-b",
+    ]
+    assert [download[1]["repo_id"] for download in downloads] == [
+        "alice/model-a",
+        "alice/model-b",
+    ]
+
+
 def test_sitecustomize_skips_sandbox_space_registration(monkeypatch):
     import huggingface_hub as hub
     from huggingface_hub import HfApi

tests/unit/test_session_manager_persistence.py

@@ -425,32 +425,9 @@ async def test_create_session_schedules_cpu_sandbox_preload():
 
         assert scheduled == [session_id]
         assert session_id in manager.sessions
-    finally:
-        stop.set()
-        await _cancel_runtime_tasks(manager)
-
-
-@pytest.mark.asyncio
-async def test_create_session_starts_hub_artifact_collection(monkeypatch):
-    manager = _manager_with_store(NoopSessionStore())
-    manager.enable_hub_artifact_collections = True
-    stop = _install_fake_runtime(manager)
-    started: list[tuple[str, str]] = []
-
-    def fake_start_session_artifact_collection_task(session, **kwargs):
-        started.append((session.session_id, kwargs["token"]))
-        return None
-
-    monkeypatch.setattr(
-        "session_manager.start_session_artifact_collection_task",
-        fake_start_session_artifact_collection_task,
-    )
-    manager._start_cpu_sandbox_preload = lambda _: None  # type: ignore[method-assign]
-
-    try:
-        session_id = await manager.create_session(user_id="owner", hf_token="token")
-
-        assert started == [(session_id, "token")]
+        runtime_session = manager.sessions[session_id].session
+        assert not hasattr(runtime_session, "_ml_intern_artifact_collection_task")
+        assert not hasattr(runtime_session, "_ml_intern_artifact_collection_slug")
     finally:
         stop.set()
         await _cancel_runtime_tasks(manager)
@@ -475,37 +452,8 @@ async def test_lazy_restore_schedules_cpu_sandbox_preload():
         assert restored is not None
         assert scheduled == ["persisted-session"]
         assert "persisted-session" in manager.sessions
-    finally:
-        stop.set()
-        await _cancel_runtime_tasks(manager)
-
-
-@pytest.mark.asyncio
-async def test_lazy_restore_starts_hub_artifact_collection(monkeypatch):
-    manager = _manager_with_store(RestoreStore())
-    manager.enable_hub_artifact_collections = True
-    stop = _install_fake_runtime(manager)
-    started: list[tuple[str, str]] = []
-
-    def fake_start_session_artifact_collection_task(session, **kwargs):
-        started.append((session.session_id, kwargs["token"]))
-        return None
-
-    monkeypatch.setattr(
-        "session_manager.start_session_artifact_collection_task",
-        fake_start_session_artifact_collection_task,
-    )
-    manager._start_cpu_sandbox_preload = lambda _: None  # type: ignore[method-assign]
-
-    try:
-        restored = await manager.ensure_session_loaded(
-            "persisted-session",
-            user_id="owner",
-            hf_token="token",
-        )
-
-        assert restored is not None
-        assert started == [("persisted-session", "token")]
+        assert not hasattr(restored.session, "_ml_intern_artifact_collection_task")
+        assert not hasattr(restored.session, "_ml_intern_artifact_collection_slug")
     finally:
         stop.set()
         await _cancel_runtime_tasks(manager)