Skip to content

feat: Implement zero-copy AsyncVectorEnv via shared memory for parallel execution - #1550

Open
shivansh023023 wants to merge 4 commits into
google-deepmind:masterfrom
shivansh023023:feature-async-vector-env
Open

feat: Implement zero-copy AsyncVectorEnv via shared memory for parallel execution#1550
shivansh023023 wants to merge 4 commits into
google-deepmind:masterfrom
shivansh023023:feature-async-vector-env

Conversation

@shivansh023023

Copy link
Copy Markdown
Contributor

Pull Request: feat: Implement zero-copy AsyncVectorEnv using shared memory

Description

This PR introduces AsyncVectorEnv to open_spiel/python/vector_env.py, enabling true parallel execution of environments across multiple CPU cores.

Key Architectural Details

  • Zero-Copy Shared Memory: To avoid severe Inter-Process Communication (IPC) pickling bottlenecks when transmitting large observation tensors (e.g., in Chess or Backgammon), this implementation utilizes multiprocessing.sharedctypes.RawArray.
  • Direct Memory Writing: Worker processes instantiate the C++ environments natively and write time_step.observations["info_state"] directly into pre-allocated shared NumPy views.
  • Lightweight Pipes: The multiprocessing.Pipe is strictly reserved for small control payloads (rewards, done flags, step types), completely eliminating the need to serialize massive multi-dimensional arrays.
  • Performance: In a 16-worker benchmark on chess running 1,000 steps, this shared memory architecture achieved a ~2.87x speedup over the synchronous SyncVectorEnv, fully bypassing the IPC overhead that typically throttles complex state transmission.

@lanctot

lanctot commented May 30, 2026

Copy link
Copy Markdown
Collaborator

Thanks!

@lanctot
lanctot self-requested a review July 11, 2026 13:11

@lanctot lanctot left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Can you add a test that shows it working on a simple game?
  • Can you add the test to the python tests in python/CMakeLists.txt

@lanctot lanctot added the waiting Waiting to hear back from contributor (tests failed or thread reply / code update required) label Jul 11, 2026
@shivansh023023

Copy link
Copy Markdown
Contributor Author
  • Can you add a test that shows it working on a simple game?
  • Can you add the test to the python tests in python/CMakeLists.txt

Done !

@lanctot lanctot removed the waiting Waiting to hear back from contributor (tests failed or thread reply / code update required) label Jul 12, 2026
@lanctot

lanctot commented Jul 12, 2026

Copy link
Copy Markdown
Collaborator

The tests failed. Can you take a look? Did they pass locally?

Comment thread open_spiel/python/vector_env.py Outdated
remote.recv()
for remote in self.remotes:
remote.send(("close", None))
for p in self.processes:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a race condition here?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(lines 279-280)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @lanctot there was indeed a potential race condition during shutdown if a worker was stuck on an IPC boundary when close() was triggered.

I've pushed a fix that handles process termination safely. I also resolved an underlying multi-player data alignment bug where info_state arrays were hitting dimension mismatches when unpacked at the master process level.

@shivansh023023

Copy link
Copy Markdown
Contributor Author

@lanctot i have implemented some chages , can you re run the checks

@shivansh023023
shivansh023023 requested a review from lanctot July 14, 2026 11:48
@shivansh023023

Copy link
Copy Markdown
Contributor Author

hi @lanctot all the checks are passing now !

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants