mirror of
https://github.com/ROCm/jax.git
synced 2025-04-29 01:26:07 +00:00
77797f434d
This change adds an experimental API `jax.experimental.colocated_python`. The ultimate goal of this API is to provide a runtime-agnostic way to wrap a Python code that runs close to (or on) accelerator hosts. Multi-controller JAX can trivially achieve this colocated Python code execution today, while single-controller JAX needed its own solution for distributed Python code execution, which creates fragmentation of the user code for these two runtime architectures. `colocated_python` is an attempt to define a single device model and portable API to allow the user to write a single code once that can run on both runtime architectures. This change includes an implementation of the function API portion of `jax.experimental.colocated_python`. A (stateful) object API will be added separately. Also there will be a separate change that expresses serialized functions as an IFRT `CustomCallProgram`. It is currently in an early development stage. Please proceed with a caution when using the API. PiperOrigin-RevId: 690705899
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
# Copyright 2024 The JAX Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Backend for colocated_python.func."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from typing import Sequence
|
|
|
|
import jax
|
|
|
|
|
|
class _ResultStore:
|
|
"""Temporarily stores results from synchronous execution of functions."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._storage: dict[int, Sequence[jax.Array]] = {}
|
|
|
|
def push(self, uid: int, out: Sequence[jax.Array]) -> None:
|
|
with self._lock:
|
|
if uid in self._storage:
|
|
raise ValueError(f"uid {uid} already exists")
|
|
self._storage[uid] = out
|
|
|
|
def pop(self, uid: int) -> Sequence[jax.Array]:
|
|
with self._lock:
|
|
if uid not in self._storage:
|
|
raise ValueError(f"uid {uid} does not exist")
|
|
return self._storage.pop(uid)
|
|
|
|
|
|
SINGLETON_RESULT_STORE = _ResultStore()
|