feat: parallel sample processing

yaak-ai · Feb 27, 2025 · be1b330 · be1b330
1 parent 6eb6844
commit be1b330
Show file tree

Hide file tree

Showing 32 changed files with 1,092 additions and 852 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -12,6 +12,9 @@ on:
 jobs:
   ci:
     runs-on: ubuntu-latest
+    env:
+      RBYTE_SAMPLES_PARALLEL: 0
+
     steps:
       - name: setup ssh
         uses: webfactory/[email protected]
@@ -48,20 +51,9 @@ jobs:
           enable-cache: true
           cache-dependency-glob: "**/pyproject.toml"
 
-      - name: sync
-        run: just sync
-
-      - name: build protos
-        run: just build-protos
-
-      - name: format
-        run: just format --check
-
-      - name: lint
-        run: just lint
-
-      - name: typecheck
-        run: just typecheck
-
-      - name: test
-        run: just test
+      - run: just sync
+      - run: just build
+      - run: just format --check
+      - run: just lint
+      - run: just typecheck
+      - run: just test
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -11,14 +11,14 @@ repos:
       - id: pyupgrade
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.5
+    rev: v0.9.7
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.26.0
+    rev: 1.28.0
     hooks:
       - id: basedpyright
 

diff --git a/config/_templates/dataset/carla.yaml b/config/_templates/dataset/carla.yaml
@@ -1,52 +1,52 @@
 #@yaml/text-templated-strings
 
 #@ drives = [
-#@     'Accident',
-#@     'Accident_2',
-#@     'AccidentTwoWays',
-#@     'BlockedIntersection',
-#@     'ConstructionObstacle',
-#@     'ConstructionObstacleTwoWays',
-#@     'ControlLoss',
-#@     'DynamicObjectCrossing_right_container',
-#@     'DynamicObjectCrossing_right_haybalelb',
-#@     'EnterActorFlow_fast',
-#@     'EnterActorFlow_slow',
-#@     'HardBreakRoute',
-#@     'HazardAtSideLane',
-#@     'HazardAtSideLaneTwoWays',
-#@     'HighwayCutIn',
-#@     'HighwayExit',
-#@     'InterurbanActorFlow_fast',
-#@     'InterurbanActorFlow_slow',
-#@     'InterurbanAdvancedActorFlow_fast',
-#@     'InterurbanAdvancedActorFlow_slow',
-#@     'InterurbanAdvancedActorFlow_slow_2',
-#@     'InvadingTurn_less_curve',
-#@     'InvadingTurn_much_curve',
-#@     'MergerIntoSlowTraffic',
-#@     'MergerIntoSlowTrafficV2',
-#@     'NonSignalizedJunctionLeftTurn_fast',
-#@     'NonSignalizedJunctionLeftTurn_slow',
-#@     'NonSignalizedJunctionRightTurn_fast',
-#@     'NonSignalizedJunctionRightTurn_slow',
-#@     'OppositeVehicleRunningRedLight',
-#@     'OppositeVehicleTakingPriority',
-#@     'ParkedObstacle',
-#@     'ParkedObstacleTwoWays',
-#@     'ParkingCrossingPedestrian',
-#@     'ParkingCutIn',
-#@     'ParkingExit',
-#@     'PedestrianCrossing',
-#@     'SignalizedJunctionLeftTurn_fast',
-#@     'SignalizedJunctionLeftTurn_slow',
-#@     'SignalizedJunctionRightTurn_fast',
-#@     'SignalizedJunctionRightTurn_slow',
-#@     'StaticCutIn',
-#@     'VehicleOpensDoorTwoWays',
-#@     'VehicleTurningRoutePedestrian',
+#!     'Accident',
+#!     'Accident_2',
+#!     'AccidentTwoWays',
+#!     'BlockedIntersection',
+#!     'ConstructionObstacle',
+#!     'ConstructionObstacleTwoWays',
+#!     'ControlLoss',
+#!     'DynamicObjectCrossing_right_container',
+#!     'DynamicObjectCrossing_right_haybalelb',
+#!     'EnterActorFlow_fast',
+#!     'EnterActorFlow_slow',
+#!     'HardBreakRoute',
+#!     'HazardAtSideLane',
+#!     'HazardAtSideLaneTwoWays',
+#!     'HighwayCutIn',
+#!     'HighwayExit',
+#!     'InterurbanActorFlow_fast',
+#!     'InterurbanActorFlow_slow',
+#!     'InterurbanAdvancedActorFlow_fast',
+#!     'InterurbanAdvancedActorFlow_slow',
+#!     'InterurbanAdvancedActorFlow_slow_2',
+#!     'InvadingTurn_less_curve',
+#!     'InvadingTurn_much_curve',
+#!     'MergerIntoSlowTraffic',
+#!     'MergerIntoSlowTrafficV2',
+#!     'NonSignalizedJunctionLeftTurn_fast',
+#!     'NonSignalizedJunctionLeftTurn_slow',
+#!     'NonSignalizedJunctionRightTurn_fast',
+#!     'NonSignalizedJunctionRightTurn_slow',
+#!     'OppositeVehicleRunningRedLight',
+#!     'OppositeVehicleTakingPriority',
+#!     'ParkedObstacle',
+#!     'ParkedObstacleTwoWays',
+#!     'ParkingCrossingPedestrian',
+#!     'ParkingCutIn',
+#!     'ParkingExit',
+#!     'PedestrianCrossing',
+#!     'SignalizedJunctionLeftTurn_fast',
+#!     'SignalizedJunctionLeftTurn_slow',
+#!     'SignalizedJunctionRightTurn_fast',
+#!     'SignalizedJunctionRightTurn_slow',
+#!     'StaticCutIn',
+#!     'VehicleOpensDoorTwoWays',
+#!     'VehicleTurningRoutePedestrian',
 #@     'VehicleTurningRoute_right',
-#@     'YieldToEmergencyVehicle',
+#!     'YieldToEmergencyVehicle',
 #@ ]
 
 #@ cameras = [
@@ -56,75 +56,85 @@
 _target_: rbyte.Dataset
 _recursive_: false
 _convert_: all
-inputs:
+sources:
   #@ for input_id in drives:
   (@=input_id@):
-    sources:
-      #@ for source_id in cameras:
-      (@=source_id@):
-        index_column: _idx_
-        source:
-          _target_: rbyte.io.PathTensorSource
-          path: "${data_dir}/(@=input_id@)/frames/(@=source_id@).defish.mp4/576x324/{:09d}.jpg"
-          decoder:
-            _target_: simplejpeg.decode_jpeg
-            _partial_: true
-            colorspace: rgb
-            fastdct: true
-            fastupsample: true
+    #@ for source_id in cameras:
+    (@=source_id@):
+      index_column: _idx_
+      source:
+        _target_: rbyte.io.PathTensorSource
+        path: "${data_dir}/(@=input_id@)/frames/(@=source_id@).defish.mp4/576x324/{:09d}.jpg"
+        decoder:
+          _target_: simplejpeg.decode_jpeg
+          _partial_: true
+          colorspace: rgb
+          fastdct: true
+          fastupsample: true
 
-      #@ end
+    #@ end
+  #@ end
 
-    samples:
-      pipeline:
-        _target_: pipefunc.Pipeline
-        validate_type_annotations: false
-        functions:
-          - _target_: pipefunc.PipeFunc
-            bound:
-              path: ${data_dir}/(@=input_id@)/ego_logs.json
-            output_name: ego_logs
-            func:
-              _target_: rbyte.io.JsonDataFrameBuilder
-              fields:
-                records:
-                  control.brake:
-                  control.throttle:
-                  control.steer:
-                  state.velocity.value:
-                  state.acceleration.value:
+samples:
+  inputs:
+    #@ for input_id in drives:
+    (@=input_id@):
+      ego_logs_path: ${data_dir}/(@=input_id@)/ego_logs.json
+    #@ end
 
-          - _target_: pipefunc.PipeFunc
-            renames:
-              input: ego_logs
-            output_name: data
-            func:
-              _target_: rbyte.io.DataFrameConcater
-              method: vertical
+  pipeline:
+    _target_: pipefunc.Pipeline
+    validate_type_annotations: false
+    functions:
+      - _target_: pipefunc.PipeFunc
+        renames:
+          path: ego_logs_path
+        output_name: ego_logs
+        mapspec: "ego_logs_path[i] -> ego_logs[i]"
+        func:
+          _target_: rbyte.io.JsonDataFrameBuilder
+          fields:
+            records:
+              control.brake:
+              control.throttle:
+              control.steer:
+              state.velocity.value:
+              state.acceleration.value:
 
-          - _target_: pipefunc.PipeFunc
-            renames:
-              input: data
-            output_name: data_resampled
-            func:
-              _target_: rbyte.io.DataFrameFpsResampler
-              fps_in: 20
-              fps_out: 30
+      - _target_: pipefunc.PipeFunc
+        renames:
+          input: ego_logs
+        output_name: data
+        mapspec: "ego_logs[i] -> data[i]"
+        func:
+          _target_: rbyte.io.DataFrameConcater
+          method: vertical
 
-          - _target_: pipefunc.PipeFunc
-            renames:
-              input: data_resampled
-            output_name: data_indexed
-            func:
-              _target_: rbyte.io.DataFrameIndexer
-              name: _idx_
+      - _target_: pipefunc.PipeFunc
+        renames:
+          input: data
+        output_name: resampled
+        mapspec: "data[i] -> resampled[i]"
+        func:
+          _target_: rbyte.io.DataFrameFpsResampler
+          fps_in: 20
+          fps_out: 30
 
-          - _target_: pipefunc.PipeFunc
-            renames:
-              input: data_indexed
-            output_name: data_filtered
-            func:
-              _target_: rbyte.io.DataFrameFilter
-              predicate: |
-                `control.throttle` > 0.5
-  #@ end
+      - _target_: pipefunc.PipeFunc
+        renames:
+          input: resampled
+        output_name: indexed
+        mapspec: "resampled[i] -> indexed[i]"
+        func:
+          _target_: rbyte.io.DataFrameIndexer
+          name: _idx_
+
+      - _target_: pipefunc.PipeFunc
+        renames:
+          input: indexed
+        output_name: samples
+        mapspec: "indexed[i] -> samples[i]"
+        func:
+          _target_: rbyte.io.DataFrameFilter
+          predicate: |
+            `control.throttle` > 0.5