From aa563dbe3662bb236ed80a3bdd65b37dec7c2eeb Mon Sep 17 00:00:00 2001 From: Sax Authors Date: Fri, 9 Aug 2024 09:23:39 -0700 Subject: [PATCH] Query and forward parameter overrides. PiperOrigin-RevId: 661290405 Change-Id: I4c51a9f542a00a94424fedf2369128aa6d073f1a --- saxml/server/BUILD | 3 --- saxml/server/model_service_main.py | 9 +-------- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/saxml/server/BUILD b/saxml/server/BUILD index 259b7f76..08a2fd20 100644 --- a/saxml/server/BUILD +++ b/saxml/server/BUILD @@ -312,9 +312,6 @@ pytype_strict_library( ":model_service_base", ":servable_model_registry", ":spmd_backend", - "//saxml/client/python:sax", - "//saxml/protobuf:admin_py_pb2", - "//saxml/protobuf:admin_py_pb2_grpc", "//saxml/protobuf:modelet_py_pb2", "//saxml/protobuf:modelet_py_pb2_grpc", "//saxml/server/jax:jax_spmd_backend", diff --git a/saxml/server/model_service_main.py b/saxml/server/model_service_main.py index 82b62218..3a4d2e02 100644 --- a/saxml/server/model_service_main.py +++ b/saxml/server/model_service_main.py @@ -22,7 +22,6 @@ import grpc import jax from jax.experimental.compilation_cache import compilation_cache -from saxml.client.python import sax from saxml.protobuf import modelet_pb2 from saxml.protobuf import modelet_pb2_grpc from saxml.server import model_service_base @@ -125,11 +124,6 @@ def _load_static_model( logging.info( 'Loading key %s, model %s, checkpoint %s.', model_key, model, checkpoint ) - # Get overrides that might have been provided via 'saxutil publish' and apply - # them. - overrides = sax.ListDetail(model_key).overrides - logging.info('Got overrides: %s', overrides) - if channel_creds is None: channel = grpc.insecure_channel(f'localhost:{port}') else: @@ -138,8 +132,7 @@ def _load_static_model( grpc.channel_ready_future(channel).result(timeout=10) stub = modelet_pb2_grpc.ModeletStub(channel) req = modelet_pb2.LoadRequest( - model_key=model_key, model_path=model, checkpoint_path=checkpoint, - overrides=overrides, + model_key=model_key, model_path=model, checkpoint_path=checkpoint ) try: stub.Load(req)