tvalentyn commented on issue #40469:
URL: https://github.com/apache/beam/issues/40469#issuecomment-6067905387
AI Repro
```
import os
import sys
import apache_beam as beam
from apache_beam.ml.inference import base
class Model:
def predict(self, x):
return x
class Handler(base.ModelHandler):
def load_model(self):
print(f'Loading model in pid {os.getpid()}', flush=True)
return Model()
def run_inference(self, batch, model, inference_args=None):
print(f'Inference on a batch of {len(batch)}', flush=True)
return [model.predict(x) for x in batch]
def batch_elements_kwargs(self):
return {'min_batch_size': 2, 'max_batch_size': 2}
if __name__ == '__main__':
use_model_manager = sys.argv[1:] != ['--no_model_manager']
with beam.Pipeline() as p:
model_metadata = p | 'ModelMetadata' >> beam.Create(
[base.ModelMetadata(model_id='model_v2', model_name='v2')])
_ = (
p
| beam.Create(range(8)) # 4 batches
| base.RunInference(
Handler(),
use_model_manager=use_model_manager,
model_metadata_pcoll=model_metadata))
```
With use_model_manager=True on the Prism runner, every batch loads the model
in a new process:
```
Loading model in pid 1348387
Inference on a batch of 2
Loading model in pid 1348560
Inference on a batch of 2
Loading model in pid 1348645
Inference on a batch of 1
Loading model in pid 1348741
Inference on a batch of 1
Loading model in pid 1348857
Inference on a batch of 1
Loading model in pid 1348938
Inference on a batch of 1
```
With --no_model_manager, the model is loaded once for all 4 batches.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
