Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

return model id in registering remote model #1329

Merged
merged 1 commit into from
Sep 13, 2023
Merged

return model id in registering remote model #1329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions ...on/src/main/java/org/opensearch/ml/common/transport/register/MLRegisterModelResponse.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,33 +18,46 @@
@Getter
public class MLRegisterModelResponse extends ActionResponse implements ToXContentObject {
public static final String TASK_ID_FIELD = "task_id";
public static final String MODEL_ID_FIELD = "model_id";
public static final String STATUS_FIELD = "status";

private String taskId;
private String status;
private String modelId;

public MLRegisterModelResponse(StreamInput in) throws IOException {
super(in);
this.taskId = in.readString();
this.status = in.readString();
this.modelId = in.readOptionalString();
}
Comment on lines 30 to 33
Copy link
Member

Choose a reason for hiding this comment

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

Feels like a good improvement for remote models, do we really have a taskId which is doing real work for remote models?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Yes, we have a actual task id. It will do some validation work

Copy link
Member

Choose a reason for hiding this comment

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

Curious what kind of validations ?


public MLRegisterModelResponse(String taskId, String status) {
this.taskId = taskId;
this.status= status;
}

public MLRegisterModelResponse(String taskId, String status, String modelId) {
this.taskId = taskId;
this.status= status;
this.modelId = modelId;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeString(taskId);
out.writeString(status);
out.writeOptionalString(modelId);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
builder.startObject();
builder.field(TASK_ID_FIELD, taskId);
builder.field(STATUS_FIELD, status);
if (modelId != null) {
builder.field(MODEL_ID_FIELD, modelId);
}
builder.endObject();
return builder;
}
Expand Down
19 changes: 18 additions & 1 deletion ...rc/test/java/org/opensearch/ml/common/transport/register/MLRegisterModelResponseTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,24 +16,27 @@ public class MLRegisterModelResponseTest {

private String taskId;
private String status;
private String modelId;

@Before
public void setUp() throws Exception {
taskId = "test_id";
status = "test";
modelId = "model_id";
}

@Test
public void writeTo_Success() throws IOException {
// Setup
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
MLRegisterModelResponse response = new MLRegisterModelResponse(taskId, status);
MLRegisterModelResponse response = new MLRegisterModelResponse(taskId, status, modelId);
// Run the test
response.writeTo(bytesStreamOutput);
MLRegisterModelResponse parsedResponse = new MLRegisterModelResponse(bytesStreamOutput.bytes().streamInput());
// Verify the results
assertEquals(response.getTaskId(), parsedResponse.getTaskId());
assertEquals(response.getStatus(), parsedResponse.getStatus());
assertEquals(response.getModelId(), parsedResponse.getModelId());
}

@Test
Expand All @@ -49,4 +52,18 @@ public void testToXContent() throws IOException {
assertEquals("{\"task_id\":\"test_id\"," +
"\"status\":\"test\"}", jsonStr);
}

@Test
public void testToXContent_withModelId() throws IOException {
// Setup
MLRegisterModelResponse response = new MLRegisterModelResponse(taskId, status, modelId);
// Run the test
XContentBuilder builder = XContentFactory.jsonBuilder();
response.toXContent(builder, ToXContent.EMPTY_PARAMS);
assertNotNull(builder);
String jsonStr = builder.toString();
// Verify the results
assertEquals("{\"task_id\":\"test_id\"," +
"\"status\":\"test\"," + "\"model_id\":\"model_id\"}", jsonStr);
}
}
5 changes: 3 additions & 2 deletions plugin/src/main/java/org/opensearch/ml/action/register/TransportRegisterModelAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ private void registerModel(MLRegisterModelInput registerModelInput, ActionListen
throw new IllegalArgumentException("URL can't match trusted url regex");
}
}
System.out.println("registering the model");
Copy link
Collaborator

Choose a reason for hiding this comment

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

Are we keeping this?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's avoid using system out print

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes I guarante these redundant lines will be removed in my next PR.

boolean isAsync = registerModelInput.getFunctionName() != FunctionName.REMOTE;
MLTask mlTask = MLTask
.builder()
Expand All @@ -249,8 +250,8 @@ private void registerModel(MLRegisterModelInput registerModelInput, ActionListen
mlTaskManager.createMLTask(mlTask, ActionListener.wrap(response -> {
String taskId = response.getId();
mlTask.setTaskId(taskId);
mlModelManager.registerMLModel(registerModelInput, mlTask);
listener.onResponse(new MLRegisterModelResponse(taskId, MLTaskState.CREATED.name()));
System.out.println("mlModelManager calls registerMLRemoteModel");
Copy link
Member

@saratvemulapalli saratvemulapalli Sep 13, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

Nit: Probably unintended?
Same at L236

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Good catch! I think I forgot to delete. They will be cleaned up for sure.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Are we keeping this line?

mlModelManager.registerMLRemoteModel(registerModelInput, mlTask, listener);
Copy link
Collaborator

Choose a reason for hiding this comment

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

Is this the expectation that, only remote model registration will be asynchronous? If yes, can we please add a comment here also?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Just the opposite, remote model registration is synchronous only so the model_id can be returned right away. Comments will be added in the next PR.

}, e -> {
logException("Failed to register model", e, log);
listener.onFailure(e);
Expand Down
Loading
Loading