Skip to content

Commit 901a5e7

Browse files
dmitshdmitsh
authored
implement stub for Nebius provider (#113)
Signed-off-by: Dmitry Shmulevich <[email protected]>
1 parent 8458bdc commit 901a5e7

File tree

5 files changed

+194
-0
lines changed

5 files changed

+194
-0
lines changed

‎go.mod‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ require (
1616
github.com/googleapis/gax-go/v2 v2.14.1
1717
github.com/hashicorp/golang-lru v1.0.2
1818
github.com/mitchellh/mapstructure v1.5.0
19+
github.com/nebius/gosdk v0.0.0-20250527091023-fab86570dc7c
1920
github.com/oklog/run v1.1.0
2021
github.com/oracle/oci-go-sdk/v65 v65.88.0
2122
github.com/prometheus/client_golang v1.20.5
@@ -32,6 +33,7 @@ require (
3233
)
3334

3435
require (
36+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.31.0-20231030212536-12f9cba37c9d.2 // indirect
3537
cloud.google.com/go/auth v0.16.1 // indirect
3638
cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect
3739
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.23 // indirect
@@ -44,6 +46,7 @@ require (
4446
github.com/aws/aws-sdk-go-v2/service/sts v1.32.3 // indirect
4547
github.com/aws/smithy-go v1.22.0 // indirect
4648
github.com/beorn7/perks v1.0.1 // indirect
49+
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
4750
github.com/cespare/xxhash/v2 v2.3.0 // indirect
4851
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
4952
github.com/emicklei/go-restful/v3 v3.12.1 // indirect
@@ -59,12 +62,14 @@ require (
5962
github.com/go-playground/universal-translator v0.18.1 // indirect
6063
github.com/gofrs/flock v0.12.1 // indirect
6164
github.com/gogo/protobuf v1.3.2 // indirect
65+
github.com/golang-jwt/jwt/v4 v4.5.1 // indirect
6266
github.com/golang/protobuf v1.5.4 // indirect
6367
github.com/google/gnostic-models v0.6.8 // indirect
6468
github.com/google/go-cmp v0.7.0 // indirect
6569
github.com/google/gofuzz v1.2.0 // indirect
6670
github.com/google/s2a-go v0.1.9 // indirect
6771
github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect
72+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.1.0 // indirect
6873
github.com/josharian/intern v1.0.0 // indirect
6974
github.com/json-iterator/go v1.1.12 // indirect
7075
github.com/klauspost/compress v1.17.11 // indirect
@@ -87,6 +92,7 @@ require (
8792
golang.org/x/crypto v0.37.0 // indirect
8893
golang.org/x/net v0.39.0 // indirect
8994
golang.org/x/oauth2 v0.30.0 // indirect
95+
golang.org/x/sync v0.14.0 // indirect
9096
golang.org/x/sys v0.33.0 // indirect
9197
golang.org/x/term v0.31.0 // indirect
9298
golang.org/x/text v0.24.0 // indirect

‎go.sum‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.31.0-20231030212536-12f9cba37c9d.2 h1:m8rKyv88R8ZIR1549RMXckZ4FZJGxrq/7aRYl6U3WHc=
2+
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.31.0-20231030212536-12f9cba37c9d.2/go.mod h1:xafc+XIsTxTy76GJQ1TKgvJWsSugFBqMaN27WhUblew=
13
cloud.google.com/go v0.121.0 h1:pgfwva8nGw7vivjZiRfrmglGWiCJBP+0OmDpenG/Fwg=
24
cloud.google.com/go v0.121.0/go.mod h1:rS7Kytwheu/y9buoDmu5EIpMMCI4Mb8ND4aeN4Vwj7Q=
35
cloud.google.com/go/auth v0.16.1 h1:XrXauHMd30LhQYVRHLGvJiYeczweKQXZxsTbV9TiguU=
@@ -36,6 +38,8 @@ github.com/aws/smithy-go v1.22.0 h1:uunKnWlcoL3zO7q+gG2Pk53joueEOsnNB28QdMsmiMM=
3638
github.com/aws/smithy-go v1.22.0/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
3739
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
3840
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
41+
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
42+
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
3943
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
4044
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
4145
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -76,10 +80,14 @@ github.com/gofrs/flock v0.12.1 h1:MTLVXXHf8ekldpJk3AKicLij9MdwOWkZ+a/jHHZby9E=
7680
github.com/gofrs/flock v0.12.1/go.mod h1:9zxTsyu5xtJ9DK+1tFZyibEV7y3uwDxPPfbxeeHCoD0=
7781
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
7882
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
83+
github.com/golang-jwt/jwt/v4 v4.5.1 h1:JdqV9zKUdtaa9gdPlywC3aeoEsR681PlKC+4F5gQgeo=
84+
github.com/golang-jwt/jwt/v4 v4.5.1/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
85+
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
7986
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
8087
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
8188
github.com/google/gnostic-models v0.6.8 h1:yo/ABAfM5IMRsS1VnXjTBvUb61tFIHozhlYvRgGre9I=
8289
github.com/google/gnostic-models v0.6.8/go.mod h1:5n7qKqH0f5wFt+aWF8CW6pZLLNOfYuF5OpfBSENuI8U=
90+
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
8391
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
8492
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
8593
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
@@ -96,6 +104,8 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU
96104
github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA=
97105
github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q=
98106
github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
107+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.1.0 h1:pRhl55Yx1eC7BZ1N+BBWwnKaMyD8uC+34TLdndZMAKk=
108+
github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.1.0/go.mod h1:XKMd7iuf/RGPSMJ/U4HP0zS2Z9Fh8Ps9a+6X26m/tmI=
99109
github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iPY6p1c=
100110
github.com/hashicorp/golang-lru v1.0.2/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
101111
github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
@@ -125,6 +135,8 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
125135
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
126136
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
127137
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
138+
github.com/nebius/gosdk v0.0.0-20250527091023-fab86570dc7c h1:UgUPIOXmD7eWJLIvWQoMLgBtes370GKCPBPCm57nZcc=
139+
github.com/nebius/gosdk v0.0.0-20250527091023-fab86570dc7c/go.mod h1:d0iaJFMQofs+hV0K25+8wVvMHdNxPmxMbIYfhOpZAt8=
128140
github.com/oklog/run v1.1.0 h1:GEenZ1cK0+q0+wsJew9qUg/DyD8k3JzYsZAi5gYi2mA=
129141
github.com/oklog/run v1.1.0/go.mod h1:sVPdnTZT1zYwAJeCMu2Th4T21pA3FPOQRfWjQlk7DVU=
130142
github.com/onsi/ginkgo/v2 v2.19.0 h1:9Cnnf7UHo57Hy3k6/m5k3dRfGTMXGvxhHFvkDTCTpvA=
@@ -240,6 +252,8 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20250428153025-10db94c68c34 h1:
240252
google.golang.org/genproto/googleapis/rpc v0.0.0-20250428153025-10db94c68c34/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
241253
google.golang.org/grpc v1.72.0 h1:S7UkcVa60b5AAQTaO6ZKamFp1zMZSU0fGDK2WZLbBnM=
242254
google.golang.org/grpc v1.72.0/go.mod h1:wH5Aktxcg25y1I3w7H69nHfXdOG3UiadoBtjh3izSDM=
255+
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
256+
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
243257
google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
244258
google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
245259
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

‎pkg/providers/nebius/instance_topology.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Copyright 2025 NVIDIA CORPORATION
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package nebius
7+
8+
import (
9+
"context"
10+
"fmt"
11+
12+
compute "github.com/nebius/gosdk/proto/nebius/compute/v1"
13+
"k8s.io/klog/v2"
14+
15+
"github.com/NVIDIA/topograph/pkg/topology"
16+
)
17+
18+
func (p *baseProvider) generateInstanceTopology(ctx context.Context, cis []topology.ComputeInstances) (*topology.ClusterTopology, error) {
19+
client, err := p.clientFactory()
20+
if err != nil {
21+
return nil, fmt.Errorf("failed to get client: %v", err)
22+
}
23+
24+
topo := topology.NewClusterTopology()
25+
26+
for _, ci := range cis {
27+
if err := p.generateRegionInstanceTopology(ctx, client, topo, &ci); err != nil {
28+
return nil, fmt.Errorf("failed to get instance topology: %v", err)
29+
}
30+
}
31+
32+
return topo, nil
33+
}
34+
35+
func (p *baseProvider) generateRegionInstanceTopology(ctx context.Context, client Client, topo *topology.ClusterTopology, ci *topology.ComputeInstances) error {
36+
if len(ci.Region) == 0 {
37+
return fmt.Errorf("must specify region")
38+
}
39+
klog.InfoS("Getting instance topology", "region", ci.Region)
40+
41+
for id, hostname := range ci.Instances {
42+
req := &compute.GetInstanceRequest{Id: id}
43+
instance, err := client.GetComputeInstance(ctx, req)
44+
if err != nil {
45+
return err
46+
}
47+
48+
ibTopology := instance.GetStatus().GetInfinibandTopologyPath()
49+
if ibTopology == nil {
50+
klog.Warningf("missing topology for node %q id %q", hostname, id)
51+
continue
52+
}
53+
54+
inst := &topology.InstanceTopology{
55+
InstanceID: id,
56+
DatacenterID: "", // TODO: set proper value from ibTopology.GetPath()
57+
SpineID: "", // TODO: set proper value from ibTopology.GetPath()
58+
BlockID: "", // TODO: set proper value from ibTopology.GetPath()
59+
}
60+
klog.Infof("Adding topology: %s", inst.String())
61+
topo.Append(inst)
62+
}
63+
return nil
64+
}

‎pkg/providers/nebius/provider.go‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/*
2+
* Copyright 2025 NVIDIA CORPORATION
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package nebius
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"os"
12+
13+
"github.com/nebius/gosdk"
14+
compute "github.com/nebius/gosdk/proto/nebius/compute/v1"
15+
services "github.com/nebius/gosdk/services/nebius/compute/v1"
16+
17+
"github.com/NVIDIA/topograph/pkg/providers"
18+
"github.com/NVIDIA/topograph/pkg/topology"
19+
)
20+
21+
const NAME = "nebius"
22+
23+
type Client interface {
24+
GetComputeInstance(context.Context, *compute.GetInstanceRequest) (*compute.Instance, error)
25+
}
26+
27+
type ClientFactory func() (Client, error)
28+
29+
type baseProvider struct {
30+
clientFactory ClientFactory
31+
}
32+
33+
type nebiusClient struct {
34+
instanceService services.InstanceService
35+
}
36+
37+
func (c *nebiusClient) GetComputeInstance(ctx context.Context, req *compute.GetInstanceRequest) (*compute.Instance, error) {
38+
return c.instanceService.Get(ctx, req)
39+
}
40+
41+
func NamedLoader() (string, providers.Loader) {
42+
return NAME, Loader
43+
}
44+
45+
func Loader(ctx context.Context, config providers.Config) (providers.Provider, error) {
46+
sdk, err := getSDK(ctx, config.Creds)
47+
if err != nil {
48+
return nil, err
49+
}
50+
51+
instanceService := sdk.Services().Compute().V1().Instance()
52+
53+
clientFactory := func() (Client, error) {
54+
return &nebiusClient{
55+
instanceService: instanceService,
56+
}, nil
57+
}
58+
59+
return New(clientFactory), nil
60+
}
61+
62+
func getSDK(ctx context.Context, creds map[string]string) (*gosdk.SDK, error) {
63+
// TODO: use credentials from payload or config to properly authenticate SDK.
64+
sdk, err := gosdk.New(
65+
ctx,
66+
gosdk.WithCredentials(
67+
gosdk.IAMToken(os.Getenv("IAM_TOKEN")),
68+
),
69+
)
70+
if err != nil {
71+
return nil, fmt.Errorf("failed to create gosdk: %v", err)
72+
}
73+
74+
return sdk, nil
75+
}
76+
77+
func (p *baseProvider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
78+
topo, err := p.generateInstanceTopology(ctx, instances)
79+
if err != nil {
80+
return nil, err
81+
}
82+
83+
return topo.ToThreeTierGraph(NAME, instances, false)
84+
}
85+
86+
type Provider struct {
87+
baseProvider
88+
}
89+
90+
func New(clientFactory ClientFactory) *Provider {
91+
return &Provider{
92+
baseProvider: baseProvider{clientFactory: clientFactory},
93+
}
94+
}
95+
96+
// Engine support
97+
98+
// Instances2NodeMap implements slurm.instanceMapper
99+
func (p *Provider) Instances2NodeMap(ctx context.Context, nodes []string) (map[string]string, error) {
100+
// TODO: implement function that returns map[instance id : hostname] for SLURM cluster given array of hostnames
101+
return nil, fmt.Errorf("not implemented")
102+
}
103+
104+
// GetComputeInstancesRegion implements slurm.instanceMapper
105+
func (p *Provider) GetComputeInstancesRegion(ctx context.Context) (string, error) {
106+
// TODO: implement function that returns region name for the current host in SLURM cluster
107+
return "", fmt.Errorf("not implemented")
108+
}

‎pkg/registry/registry.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"github.com/NVIDIA/topograph/pkg/providers/baremetal"
2727
"github.com/NVIDIA/topograph/pkg/providers/cw"
2828
"github.com/NVIDIA/topograph/pkg/providers/gcp"
29+
"github.com/NVIDIA/topograph/pkg/providers/nebius"
2930
"github.com/NVIDIA/topograph/pkg/providers/oci"
3031
provider_test "github.com/NVIDIA/topograph/pkg/providers/test"
3132
)
@@ -40,6 +41,7 @@ var Providers = providers.NewRegistry(
4041
oci.NamedLoaderAPI,
4142
oci.NamedLoaderIMDS,
4243
oci.NamedLoaderSim,
44+
nebius.NamedLoader,
4345
provider_test.NamedLoader,
4446
)
4547

0 commit comments

Comments
 (0)