Skip to content

Add cluster service and SelectCluster method #7187

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

Merged
katrogan merged 4 commits into from
Apr 10, 2026
Merged

Add cluster service and SelectCluster method #7187

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dbfff25
Add cluster service and SelectCluster method
katrogan Apr 9, 2026
416866e
Merge branch 'v2' into katrina/eng26-351-create-a-selectcluster-rpc
katrogan Apr 10, 2026
13dc4ef
regen
katrogan Apr 10, 2026
9d4268f
revert mocks
katrogan Apr 10, 2026
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
32 changes: 32 additions & 0 deletions dataproxy/service/cluster_service.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
package service

import (
"context"

"connectrpc.com/connect"

"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect"
)

type ClusterService struct {
clusterconnect.UnimplementedClusterServiceHandler
dataplaneDomain string
}

func NewClusterService(dataplaneDomain string) *ClusterService {
return &ClusterService{
dataplaneDomain: dataplaneDomain,
}
}

var _ clusterconnect.ClusterServiceHandler = (*ClusterService)(nil)

func (s *ClusterService) SelectCluster(
ctx context.Context,
req *connect.Request[cluster.SelectClusterRequest],
) (*connect.Response[cluster.SelectClusterResponse], error) {
return connect.NewResponse(&cluster.SelectClusterResponse{
ClusterEndpoint: s.dataplaneDomain,
}), nil
}
6 changes: 6 additions & 0 deletions dataproxy/setup.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"github.com/flyteorg/flyte/v2/flytestdlib/app"
"github.com/flyteorg/flyte/v2/dataproxy/config"
"github.com/flyteorg/flyte/v2/dataproxy/service"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/dataproxy/dataproxyconnect"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/task/taskconnect"
"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/trigger/triggerconnect"
Expand All @@ -30,6 +31,11 @@ func Setup(ctx context.Context, sc *app.SetupContext) error {
sc.Mux.Handle(path, handler)
logger.Infof(ctx, "Mounted DataProxyService at %s", path)

clusterSvc := service.NewClusterService(baseURL)
Copy link
Copy Markdown
Contributor

@EngHabu EngHabu Apr 9, 2026

Choose a reason for hiding this comment

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

Does this have dns:/// prefix? http://? does it matter?

Copy link
Copy Markdown
Contributor Author

@katrogan katrogan Apr 10, 2026

Choose a reason for hiding this comment

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

http but i was thinking clients can normalize it as they please? can also switch to dns if you think that makes more sense!

 $ curl -s -X POST http://localhost:30080/flyteidl2.cluster.ClusterService/SelectCluster \                                                                                            
    -H "Content-Type: application/json" \                                                                                                                                              
    -d '{"resource": {"orgId": {"name": "test"}}, "operation": "OPERATION_CREATE_UPLOAD_LOCATION"}'                                                                                    
                                                            
  {"clusterEndpoint":"http://localhost:8090"}

Copy link
Copy Markdown
Contributor

@EngHabu EngHabu Apr 10, 2026

Choose a reason for hiding this comment

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

We expect clients to be using connect, in which case I think http works?

clusterPath, clusterHandler := clusterconnect.NewClusterServiceHandler(clusterSvc)
sc.Mux.Handle(clusterPath, clusterHandler)
logger.Infof(ctx, "Mounted ClusterService at %s", clusterPath)

sc.AddReadyCheck(func(r *http.Request) error {
baseContainer := sc.DataStore.GetBaseContainerFQN(r.Context())
if baseContainer == "" {
Expand Down
41 changes: 41 additions & 0 deletions flyteidl2/cluster/payload.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
syntax = "proto3";

package flyteidl2.cluster;

import "buf/validate/validate.proto";
import "flyteidl2/app/app_definition.proto";
import "flyteidl2/common/identifier.proto";
import "flyteidl2/task/task_definition.proto";

option go_package = "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster";

message SelectClusterRequest {
oneof resource {
option (buf.validate.oneof).required = true;
flyteidl2.common.OrgIdentifier org_id = 1;
flyteidl2.common.ProjectIdentifier project_id = 2;
flyteidl2.task.TaskIdentifier task_id = 3;
flyteidl2.common.ActionIdentifier action_id = 4;
flyteidl2.common.ActionAttemptIdentifier action_attempt_id = 5;
flyteidl2.app.Identifier app_id = 6;
}

enum Operation {
OPERATION_UNSPECIFIED = 0;
OPERATION_CREATE_UPLOAD_LOCATION = 1;
OPERATION_UPLOAD_INPUTS = 2;
OPERATION_GET_ACTION_DATA = 3;
OPERATION_QUERY_RANGE_METRICS = 4;
OPERATION_CREATE_DOWNLOAD_LINK = 5;
OPERATION_TAIL_LOGS = 6;
OPERATION_GET_ACTION_ATTEMPT_METRICS = 7;
}

Operation operation = 8 [(buf.validate.field).enum = {
not_in: [0]
}];
}

message SelectClusterResponse {
string cluster_endpoint = 1;
}
17 changes: 17 additions & 0 deletions flyteidl2/cluster/service.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
syntax = "proto3";

package flyteidl2.cluster;

import "flyteidl2/cluster/payload.proto";
import "google/api/annotations.proto";

option go_package = "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster";

service ClusterService {
rpc SelectCluster(SelectClusterRequest) returns (SelectClusterResponse) {
option (google.api.http) = {
post: "/cluster/api/v1/select"
body: "*"
};
Copy link
Copy Markdown
Contributor

@EngHabu EngHabu Apr 9, 2026

Choose a reason for hiding this comment

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

No google/api annotations needed.. we don't use gateway anymore

Copy link
Copy Markdown
Contributor Author

@katrogan katrogan Apr 10, 2026

Choose a reason for hiding this comment

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

thanks, updated

}
}
113 changes: 113 additions & 0 deletions gen/go/flyteidl2/cluster/clusterconnect/service.connect.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading