[release-0.7] feat: retract executor attributes when cached data is evicted

docs/api/types.md

@@ -407,18 +407,29 @@ message ExecutorSpec {
 
 ### ExecutorAttributes
 
-Opaque locality keys published by a service instance.
+Opaque locality keys published, or retracted, by a service instance.
 
 ```protobuf
 message ExecutorAttributes {
-  repeated bytes attr = 1;
+  repeated bytes publish = 1;
+  repeated bytes retract = 2;
 }
 ```
 
-Session Manager extends the executor's existing attribute set with every
-transported value. Keys remain until the executor is removed. Keys must be
-nonempty and no larger than 256 bytes; one publication may contain at most
-1,024 distinct keys and 64 KiB after deduplication.
+Session Manager adds every `publish` key to the executor's attribute set and
+removes every `retract` key from it. Keys not retracted remain until the
+executor is removed, including across unbind and rebind. Retracting an absent
+key is harmless, and a later publication can add it again. An update with both
+lists empty changes nothing.
+
+Keys must be nonempty and no larger than 256 bytes; each list may contain at
+most 1,024 distinct keys and 64 KiB after deduplication. Session Manager
+validates the whole update first and rejects it, leaving the attribute set
+unchanged, if either list is invalid or exceeds its limit, or if a key appears
+in both lists.
+
+Publisher and Session Manager must be upgraded together: a Session Manager
+that predates `retract` ignores the field instead of removing the keys.
 
 ### ExecutorStatus
 

docs/designs/RFE275-data-aware-scheduling/FS.md

@@ -184,7 +184,9 @@ is not overloaded:
 
 ```protobuf
 message ExecutorAttributes {
-  repeated bytes attr = 1;
+  repeated bytes publish = 1;
+  // Added by #632: keys to remove from the retained set.
+  repeated bytes retract = 2;
 }
 
 message OnTaskInvokeResponse {

docs/sdk/python.md

@@ -392,10 +392,12 @@ API names such as `close`.
 
 During an App invocation, `app.session_context()` returns the active session
 context, while `app.publish_attributes(attrs)` adds opaque `bytes` keys to the
-current response. Repeated calls in one response accumulate; App publishes and
-drains the set at the response boundary. Session Manager unions every response
-into the executor's retained attribute set. Task calls can request a matching
-instance with `TaskOptions(affinity={key})`.
+current response, and `app.retract_attributes(attrs)` retracts keys that are
+no longer valid, such as a cache entry that was evicted. Repeated calls in one
+response accumulate, with the last call on a key winning; App sends and drains
+the changes at the response boundary. Session Manager adds the published keys
+to the executor's retained attribute set and removes the retracted ones. Task
+calls can request a matching instance with `TaskOptions(affinity={key})`.
 
 The returned service-side context is the core `flamepy.SessionContext`. An inner
 `@app.service()` declaration made during a service invocation can call
@@ -437,7 +439,7 @@ The repository-level check validates the `flmrun` template, App package upload/i
 | Applications | `register_application()`, `unregister_application()`, `get_application()`, `list_applications()` |
 | Services | `FlameService`, `flamepy.run()`, `flamepy.serving.Instance`, `flamepy.serving.open_session()` |
 | Objects | `put_object()`, `get_object()`, `patch_object()`, `delete_objects()`, `upload_package()`, `download_package()` |
-| App | `flamepy.app.init()`, `flamepy.app.service()`, `flamepy.app.destroy()`, `flamepy.app.session_context()`, `flamepy.app.publish_attributes()`, `ServiceInstance`, `ObjectFuture` |
+| App | `flamepy.app.init()`, `flamepy.app.service()`, `flamepy.app.destroy()`, `flamepy.app.session_context()`, `flamepy.app.publish_attributes()`, `flamepy.app.retract_attributes()`, `ServiceInstance`, `ObjectFuture` |
 
 See also:
 

docs/sdk/rust.md

@@ -213,11 +213,14 @@ impl FlameService for RawService {
 
 `self.publish()` and the runtime-provided `FlameInstance::publish()` add opaque
 locality keys to the next successful session-enter response or the next task
-response, including a failed task. Calls in one round are unioned, and the
-transported set extends the executor's retained attributes. Keys must be
-nonempty and at most 256 bytes; a round supports at most 1,024 distinct keys and
-64 KiB after deduplication. A manually constructed `FlameInstance` is not
-attached to a service publisher, so its `publish()` method returns an error.
+response, including a failed task. `retract()` withdraws keys, for example
+after a cache eviction. Calls in one round are merged with the last operation on
+a key winning, and the transported changes update the executor's retained
+attributes. Keys must be nonempty and at most 256 bytes; a round supports at
+most 1,024 distinct keys and 64 KiB after deduplication, separately for
+published and retracted keys. A manually constructed `FlameInstance` is not
+attached to a service publisher, so its `publish()` and `retract()` methods
+return an error.
 
 ## Use Object Cache
 

docs/tutorials/app-setup.md

@@ -351,10 +351,13 @@ print(result.get())
 app.destroy()
 ```
 
-App publishes accumulated keys after every task invocation, including a failed
-task. Session Manager retains previously accepted keys, and publishing nothing
-is a no-op. Keys must be nonempty and at most 256 bytes. One publication
-round supports at most 1,024 distinct keys and 64 KiB after deduplication.
+App sends the accumulated changes after every task invocation, including a
+failed task. Session Manager retains previously accepted keys, and publishing
+nothing is a no-op. When the cache drops an entry, call
+`app.retract_attributes({key})` so the executor stops advertising it; keys
+that are not retracted stay advertised across rebinding. Keys must be nonempty
+and at most 256 bytes. One round supports at most 1,024 distinct keys and
+64 KiB after deduplication, separately for published and retracted keys.
 
 App constructs class execution objects from the supplied constructor arguments
 in each executor process. It retains an object by its handle's service ID when

examples/vllm/offline/tests/test_engine.py

@@ -7,7 +7,7 @@
 import unittest
 from unittest.mock import patch
 
-from flamepy.app._context import _bind_invocation_context
+from flamepy.app._context import _bind_invocation_context, _stage_response_attributes
 from flamepy.app.runpy import FlameRunpyService
 from flamepy.app.types import ServiceContext
 from flamepy.core.service import ApplicationContext, SessionContext
@@ -80,15 +80,14 @@ def test_service_publishes_prompt_keys(self) -> None:
         with _bind_invocation_context(session) as invocation_context:
             first = execution_object.generate("first")
             second = execution_object.generate("second")
-        service.publish(invocation_context.attributes)
+        _stage_response_attributes(invocation_context.attributes)
 
         self.assertEqual(FakeLLM.creations, 1)
         self.assertEqual(first.cached_tokens, 0)
         self.assertEqual(second.cached_tokens, 0)
-        self.assertEqual(
-            set(service._take_attributes().attr),
-            {kv_key(MODEL, "first"), kv_key(MODEL, "second")},
-        )
+        attributes = service._take_attributes()
+        self.assertEqual(set(attributes.publish), {kv_key(MODEL, "first"), kv_key(MODEL, "second")})
+        self.assertEqual(list(attributes.retract), [])
 
     def test_rebinding_reuses_cached_prompt(self) -> None:
         service = FlameRunpyService()
@@ -98,24 +97,23 @@ def test_rebinding_reuses_cached_prompt(self) -> None:
         session = SessionContext(None, "session", ApplicationContext("vllm-das"))
         with _bind_invocation_context(session) as invocation_context:
             first = first_execution_object.generate("same prompt")
-        service.publish(invocation_context.attributes)
+        _stage_response_attributes(invocation_context.attributes)
         service._take_attributes()
 
         asyncio.run(service.on_session_leave())
         service._set_execution_from_context(context)
         second_execution_object = service._execution_object
         with _bind_invocation_context(session) as invocation_context:
             second = second_execution_object.generate("same prompt")
-        service.publish(invocation_context.attributes)
+        _stage_response_attributes(invocation_context.attributes)
 
         self.assertEqual(first.cached_tokens, 0)
         self.assertGreater(second.cached_tokens, 0)
         self.assertEqual(FakeLLM.creations, 1)
         self.assertIs(second_execution_object, first_execution_object)
-        self.assertEqual(
-            set(service._take_attributes().attr),
-            {kv_key(MODEL, "same prompt")},
-        )
+        attributes = service._take_attributes()
+        self.assertEqual(set(attributes.publish), {kv_key(MODEL, "same prompt")})
+        self.assertEqual(list(attributes.retract), [])
 
 
 if __name__ == "__main__":

rpc/protos/types.proto

@@ -76,9 +76,11 @@ message TaskSpec {
 }
 
 // Volatile, executor-local data advertised by a service for affinity routing.
-// A publication extends the executor's retained attribute set.
+// An update adds `publish` keys to, and removes `retract` keys from, the
+// executor's retained attribute set. A key must not appear in both lists.
 message ExecutorAttributes {
-  repeated bytes attr = 1;
+  repeated bytes publish = 1;
+  repeated bytes retract = 2;
 }
 
 message Task {

sdk/python/docs/API.md

@@ -206,8 +206,11 @@ Key classes and helpers:
   members have joined, such as a collective `init_process_group`, can use it as
   the world size. It raises `RuntimeError` outside an invocation.
 - `app.publish_attributes(attrs)` adds opaque `bytes` locality keys to the
-  current App response; repeated calls in the response accumulate. Session
-  Manager unions every response into the executor's retained set.
+  current App response; repeated calls in the response accumulate.
+- `app.retract_attributes(attrs)` retracts keys that are no longer valid,
+  such as an evicted cache entry. The last call on a key within a response wins.
+  Session Manager adds published keys to the executor's retained set and
+  removes retracted ones; retracting an absent key is harmless.
 - `init(name, fail_if_exists=False, dependencies=None, python_version=None)`
   packages and initializes the application and returns its runtime handle.
   Repeating it for the same name returns the same handle. `dependencies`

sdk/python/protos/types.proto

@@ -76,9 +76,11 @@ message TaskSpec {
 }
 
 // Volatile, executor-local data advertised by a service for affinity routing.
-// A publication extends the executor's retained attribute set.
+// An update adds `publish` keys to, and removes `retract` keys from, the
+// executor's retained attribute set. A key must not appear in both lists.
 message ExecutorAttributes {
-  repeated bytes attr = 1;
+  repeated bytes publish = 1;
+  repeated bytes retract = 2;
 }
 
 message Task {

sdk/python/src/flamepy/app/__init__.py

@@ -11,7 +11,7 @@
 limitations under the License.
 """
 
-from flamepy.app._context import InstanceContext, context, publish_attributes, session_context
+from flamepy.app._context import InstanceContext, context, publish_attributes, retract_attributes, session_context
 from flamepy.app.types import ServiceContext, ServiceRequest, ServiceResponse
 from flamepy.core import Ref, ValueRef
 
@@ -53,6 +53,7 @@
     "publish_attributes",
     "ref",
     "remote",
+    "retract_attributes",
     "select",
     "service",
     "session_context",

sdk/python/src/flamepy/app/_context.py

@@ -18,9 +18,10 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import FrozenSet, Iterable, Iterator, Optional
+from typing import Iterable, Iterator, Optional
 
-from flamepy.core.service import SessionContext, TaskContext, _validate_attributes
+from flamepy.core.service import SessionContext, TaskContext, _AttributeDelta, _AttributeOp, _validate_attributes
+from flamepy.proto.types_pb2 import ExecutorAttributes
 
 
 @dataclass(frozen=True)
@@ -44,7 +45,7 @@ class InstanceContext:
 @dataclass
 class _InvocationContext:
     session_context: SessionContext
-    attributes: set[bytes]
+    attributes: _AttributeDelta
     task_context: Optional[TaskContext] = None
 
 
@@ -57,7 +58,7 @@ class _InvocationContext:
 # ``on_task_invoke`` returns. Keeping this handoff in a ContextVar preserves the
 # invoking thread/task's attributes until that response is assembled, without
 # putting concurrent invocations through the core service's shared publisher.
-_pending_response_attributes: ContextVar[Optional[FrozenSet[bytes]]] = ContextVar(
+_pending_response_attributes: ContextVar[Optional[ExecutorAttributes]] = ContextVar(
     "flamepy_app_pending_response_attributes",
     default=None,
 )
@@ -84,15 +85,27 @@ def context() -> InstanceContext:
     )
 
 
-def publish_attributes(attributes: Iterable[bytes]) -> None:
-    """Publish executor-local affinity attributes for the current response."""
+def _current_invocation() -> _InvocationContext:
     context = _current_invocation_context.get()
     if context is None:
         raise RuntimeError("App code is not running in a Flame invocation")
-    incoming = _validate_attributes(attributes)
-    combined = _validate_attributes((*context.attributes, *incoming))
-    context.attributes.clear()
-    context.attributes.update(combined)
+    return context
+
+
+def publish_attributes(attributes: Iterable[bytes]) -> None:
+    """Publish executor-local affinity attributes for the current response."""
+    context = _current_invocation()
+    context.attributes.record(_validate_attributes(attributes), _AttributeOp.PUBLISH)
+
+
+def retract_attributes(attributes: Iterable[bytes]) -> None:
+    """Retract affinity attributes that are no longer valid, e.g. after a cache eviction.
+
+    The last of ``publish_attributes``/``retract_attributes`` on a key within
+    one invocation wins.
+    """
+    context = _current_invocation()
+    context.attributes.record(_validate_attributes(attributes), _AttributeOp.RETRACT)
 
 
 @contextmanager
@@ -103,7 +116,7 @@ def _bind_invocation_context(
     """Bind runtime capabilities while user App code is executing."""
     invocation_context = _InvocationContext(
         session_context=session,
-        attributes=set(),
+        attributes=_AttributeDelta(),
         task_context=task,
     )
     token = _current_invocation_context.set(invocation_context)
@@ -113,12 +126,12 @@ def _bind_invocation_context(
         _current_invocation_context.reset(token)
 
 
-def _stage_response_attributes(attributes: Iterable[bytes]) -> None:
+def _stage_response_attributes(attributes: _AttributeDelta) -> None:
     """Stage attributes for the current invocation's shim response."""
-    _pending_response_attributes.set(_validate_attributes(attributes))
+    _pending_response_attributes.set(attributes.take())
 
 
-def _take_response_attributes() -> Optional[FrozenSet[bytes]]:
+def _take_response_attributes() -> Optional[ExecutorAttributes]:
     """Take attributes staged for the current invocation, if any."""
     attributes = _pending_response_attributes.get()
     if attributes is not None:

sdk/python/src/flamepy/app/runpy.py

@@ -58,7 +58,7 @@ def _take_attributes(self) -> ExecutorAttributes:
         """Take attributes belonging to the invocation assembling this response."""
         attributes = _take_response_attributes()
         if attributes is not None:
-            return ExecutorAttributes(attr=list(attributes))
+            return attributes
         return super()._take_attributes()
 
     def _set_execution_from_context(self, service_context: ServiceContext) -> None: