buildgrid 0.3.3__tar.gz → 0.3.4__tar.gz

{buildgrid-0.3.3 → buildgrid-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.3.3
+Version: 0.3.4
 Summary: A remote execution service
 License: Apache License, Version 2.0
 Project-URL: Homepage, https://buildgrid.build

{buildgrid-0.3.3 → buildgrid-0.3.4}/buildgrid/server/scheduler/impl.py

@@ -2464,14 +2464,16 @@ class Scheduler:
         action_result.execution_metadata.worker_completed_timestamp.FromDatetime(job.worker_completed_timestamp or now)
         response = ExecuteResponse(result=action_result, cached_result=False, status=status)
 
-        job.result = digest_to_string(self.storage.put_message(response))
+        with instance_context(job.instance_name):
+            job.result = digest_to_string(self.storage.put_message(response))
 
         self._update_job_timestamps(session, job, action_result.execution_metadata)
 
         if self.action_cache and result and not job.do_not_cache:
             action_digest = string_to_digest(job.action_digest)
             try:
-                self.action_cache.update_action_result(action_digest, action_result)
+                with instance_context(job.instance_name):
+                    self.action_cache.update_action_result(action_digest, action_result)
                 LOGGER.debug(
                     "Stored action result in ActionCache.",
                     tags=dict(action_result=action_result, digest=action_digest),
@@ -2662,7 +2664,9 @@ class Scheduler:
         return metadata
 
     def _fetch_execution_stats(
-        self, auxiliary_metadata: RepeatedCompositeFieldContainer[ProtoAny]
+        self,
+        auxiliary_metadata: RepeatedCompositeFieldContainer[ProtoAny],
+        instance_name: str,
     ) -> ExecutionStatistics | None:
         """Fetch ExecutionStatistics from Storage
         ProtoAny[Digest] -> ProtoAny[ExecutionStatistics]
@@ -2675,7 +2679,8 @@ class Scheduler:
             try:
                 aux_metadata_any.Unpack(aux_metadata_digest)
                 # Get the blob from CAS
-                execution_stats_any = self.storage.get_message(aux_metadata_digest, ProtoAny)
+                with instance_context(instance_name):
+                    execution_stats_any = self.storage.get_message(aux_metadata_digest, ProtoAny)
                 # Get the wrapped ExecutionStatistics
                 if execution_stats_any and execution_stats_any.Is(ExecutionStatistics.DESCRIPTOR):
                     execution_stats = ExecutionStatistics()
@@ -2739,7 +2744,7 @@ class Scheduler:
         if self.metering_client is None or len(execution_metadata.auxiliary_metadata) == 0:
             return
 
-        execution_stats = self._fetch_execution_stats(execution_metadata.auxiliary_metadata)
+        execution_stats = self._fetch_execution_stats(execution_metadata.auxiliary_metadata, instance_name)
         if execution_stats is None:
             return
         usage = Usage(
@@ -2809,7 +2814,8 @@ class Scheduler:
         try:
             # BuildGrid doesn't store action_result in CAS, but if we push it as an asset
             # we need it to be accessible
-            digest = self.storage.put_message(action_result)
+            with instance_context(instance_name):
+                digest = self.storage.put_message(action_result)
 
             uri = DIGEST_URI_TEMPLATE.format(digest_hash=digest.hash)
             qualifier = {"resource_type": PROTOBUF_MEDIA_TYPE}

{buildgrid-0.3.3 → buildgrid-0.3.4}/buildgrid/server/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 
-__version__ = "0.3.3"
+__version__ = "0.3.4"

{buildgrid-0.3.3 → buildgrid-0.3.4}/buildgrid.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.3.3
+Version: 0.3.4
 Summary: A remote execution service
 License: Apache License, Version 2.0
 Project-URL: Homepage, https://buildgrid.build

{buildgrid-0.3.3 → buildgrid-0.3.4}/pyproject.toml

@@ -21,7 +21,7 @@ build-backend = 'setuptools.build_meta'
 
 [project]
 name = "buildgrid"
-version = "0.3.3"
+version = "0.3.4"
 requires-python = ">=3.11"
 description = "A remote execution service"
 readme = "README.rst"

{buildgrid-0.3.3 → buildgrid-0.3.4}/tests/test_scheduler.py

@@ -41,8 +41,10 @@ from buildgrid._protos.build.buildgrid.scheduling_pb2 import SchedulingMetadata
 from buildgrid._protos.google.rpc import code_pb2
 from buildgrid._protos.google.rpc.status_pb2 import Status
 from buildgrid.server.actioncache.caches.lru_cache import LruActionCache
+from buildgrid.server.actioncache.caches.remote_cache import RemoteActionCache
 from buildgrid.server.cas.storage import lru_memory_cache
-from buildgrid.server.cas.storage.storage_abc import create_write_session
+from buildgrid.server.cas.storage.remote import RemoteStorage
+from buildgrid.server.cas.storage.storage_abc import StorageABC, create_write_session
 from buildgrid.server.client.asset import AssetClient
 from buildgrid.server.context import instance_context
 from buildgrid.server.controller import ExecutionController
@@ -75,6 +77,8 @@ from buildgrid.server.sql.models import (
 )
 from buildgrid.server.sql.provider import SqlProvider
 from buildgrid.server.utils.digests import create_digest
+from tests.utils.action_cache import serve_cache
+from tests.utils.cas import serve_cas
 from tests.utils.fixtures import (  # noqa: F401 allow fixture inclusion
     connection_string_from_connection,
     mock_logstream_client,
@@ -136,37 +140,36 @@ def cohort_set() -> CohortSet:
     )
 
 
-PARAMS = ["action-cache", "no-action-cache"]
+USE_ACTION_CACHE = ["action-cache", "remote-action-cache", "no-action-cache"]
+USE_REMOTE_STORAGE = ["lru-storage", "remote-storage"]
 
 
-@pytest.fixture(params=PARAMS)
-def controller(
-    new_postgres, request, asset_client: AssetClient, common_props: PropertySet, cohort_set: CohortSet
-) -> Iterator[ExecutionController]:
-    use_cache = request.param
-    storage = lru_memory_cache.LRUMemoryCache(1024 * 1024)
-    connection = connection_string_from_connection(new_postgres)
-    sql_provider = SqlProvider(connection_string=connection)
+@pytest.fixture(params=USE_REMOTE_STORAGE)
+def storage(request) -> Iterator[StorageABC]:
+    if request.param == "lru-storage":
+        yield lru_memory_cache.LRUMemoryCache(1024 * 1024)
+    elif request.param == "remote-storage":
+        with serve_cas(["sql", "linux"]) as server:
+            yield RemoteStorage(server.remote, max_backoff=0)
 
-    # Create database tables
-    with sql_provider._engine.begin() as conn:
-        Base.metadata.create_all(conn)
 
-    with create_write_session(command_digest) as write_session:
-        write_session.write(command.SerializeToString())
-        storage.commit_write(command_digest, write_session)
-
-    with create_write_session(action_digest) as write_session:
-        write_session.write(action.SerializeToString())
-        storage.commit_write(action_digest, write_session)
+@pytest.fixture(params=USE_ACTION_CACHE)
+def scheduler(
+    new_postgres,
+    request,
+    asset_client: AssetClient,
+    common_props: PropertySet,
+    cohort_set: CohortSet,
+    storage: StorageABC,
+) -> Iterator[Scheduler]:
+    connection = connection_string_from_connection(new_postgres)
+    sql_provider = SqlProvider(connection_string=connection)
 
-    if use_cache == "action-cache":
-        cache = LruActionCache(storage, 50)
-        scheduler = Scheduler(
+    if request.param == "no-action-cache":
+        yield Scheduler(
             sql_provider,
             storage,
             poll_interval=0.01,
-            action_cache=cache,
             property_set=common_props,
             cohort_set=cohort_set,
             asset_client=asset_client,
@@ -174,11 +177,13 @@ def controller(
             completed_action_retention_hours=1,
             action_result_retention_hours=3,
         )
-    else:
-        scheduler = Scheduler(
+    elif request.param == "action-cache":
+        cache = LruActionCache(storage, 50)
+        yield Scheduler(
             sql_provider,
             storage,
             poll_interval=0.01,
+            action_cache=cache,
             property_set=common_props,
             cohort_set=cohort_set,
             asset_client=asset_client,
@@ -186,8 +191,41 @@ def controller(
             completed_action_retention_hours=1,
             action_result_retention_hours=3,
         )
+    elif request.param == "remote-action-cache":
+        with serve_cache(["sql", "linux"]) as cache_server:
+            cache = RemoteActionCache(cache_server.remote)
+            yield Scheduler(
+                sql_provider,
+                storage,
+                poll_interval=0.01,
+                action_cache=cache,
+                property_set=common_props,
+                cohort_set=cohort_set,
+                asset_client=asset_client,
+                queued_action_retention_hours=12,
+                completed_action_retention_hours=1,
+                action_result_retention_hours=3,
+            )
+
+
+@pytest.fixture
+def controller(
+    storage: StorageABC,
+    scheduler: Scheduler,
+) -> Iterator[ExecutionController]:
+    # Create database tables
+    with scheduler._sql._engine.begin() as conn:
+        Base.metadata.create_all(conn)
 
     with scheduler, instance_context("sql"):
+        with create_write_session(command_digest) as write_session:
+            write_session.write(command.SerializeToString())
+            storage.commit_write(command_digest, write_session)
+
+        with create_write_session(action_digest) as write_session:
+            write_session.write(action.SerializeToString())
+            storage.commit_write(action_digest, write_session)
+
         yield ExecutionController(scheduler)
 
 
@@ -882,7 +920,7 @@ def test_publish_execution_stats(controller, mock_metering_client, mock_executio
     scheduler.metering_client = mock_metering_client
 
     with instance_context("sql"):
-        scheduler.publish_execution_stats("job", "", mock_execution_metadata, "linux")
+        scheduler.publish_execution_stats("job", "sql", mock_execution_metadata, "linux")
     assert len(mock_metering_client.sink) == 2
 
 
@@ -892,7 +930,7 @@ def test_publish_execution_stats_fetch_fails_no_throw(controller, mock_metering_
     mock_metering_client.put_usage.side_effect = MeteringServiceHTTPError(503, "big bang")
 
     with instance_context("sql"):
-        scheduler.publish_execution_stats("job", "", mock_execution_metadata)
+        scheduler.publish_execution_stats("job", "sql", mock_execution_metadata)
     assert len(mock_metering_client.sink) == 0
 
 
@@ -907,7 +945,7 @@ def test_publish_execution_stats_invalid_execution_stats(controller, mock_meteri
     mock_execution_metadata.auxiliary_metadata.append(some_digest_any)
 
     with instance_context("sql"):
-        scheduler.publish_execution_stats("job", "", mock_execution_metadata, "linux")
+        scheduler.publish_execution_stats("job", "sql", mock_execution_metadata, "linux")
     assert len(mock_metering_client.sink) == 0
 
 
@@ -916,7 +954,7 @@ def test_publish_execution_stats_no_client(controller, mock_metering_client, moc
     # Just dont assign a client?? Weird test
 
     with instance_context("sql"):
-        scheduler.publish_execution_stats("job", "", mock_execution_metadata, "linux")
+        scheduler.publish_execution_stats("job", "sql", mock_execution_metadata, "linux")
     assert len(mock_metering_client.sink) == 0