Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 4 additions & 1 deletion api/dataset/v1alpha1/dataset_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ const (
DatasetTypeModelScope DatasetType = "MODEL_SCOPE"
DatasetTypeDatabase DatasetType = "DATABASE"
DatasetTypeHadoop DatasetType = "HADOOP"
DatasetTypeManual DatasetType = "MANUAL"

// must be same as apis/management-api/dataset/v1alpha1/dataset.proto
DatasetStatusPhasePending DatasetStatusPhase = "PENDING"
Expand All @@ -51,7 +52,7 @@ const (
)

type DatasetSource struct {
// +kubebuilder:validation:Enum=GIT;S3;HTTP;PVC;NFS;CONDA;REFERENCE;HUGGING_FACE;MODEL_SCOPE;DATABASE;HADOOP
// +kubebuilder:validation:Enum=GIT;S3;HTTP;PVC;NFS;CONDA;REFERENCE;HUGGING_FACE;MODEL_SCOPE;DATABASE;HADOOP;MANUAL
// +kubebuilder:validation:XValidation:rule="self == oldSelf",message="Value is immutable"
Type DatasetType `json:"type"`
// +kubebuilder:validation:Required
Expand All @@ -68,6 +69,7 @@ type DatasetSource struct {
// - MODEL_SCOPE: modelscope://<namespace>/<model>
// - DATABASE: database://<ip>:<port>
// - HADOOP: hdfs://<ip>:<port>
// - MANUAL: manual://
// +kubebuilder:validation:XValidation:rule="self == oldSelf",message="Value is immutable"
URI string `json:"uri"`
// +kubebuilder:validation:Optional
Expand All @@ -86,6 +88,7 @@ type DatasetSource struct {
// * Note: syncMode can be "sync" (default) or "copy". "sync" removes files in destination that don't exist in source, "copy" only adds/updates files without removing existing ones.
// - DATABASE: type(currently only support MySQL, other database types may be supported in the future.), host, port, dbName, tables(in the dbName), exportFormat(currently only support csv)
// - HADOOP: coreSiteXml and hdfsSiteXml, sourcePath, username
// - MANUAL:
Options map[string]string `json:"options,omitempty"`
}

Expand Down
3 changes: 3 additions & 0 deletions config/crd/bases/dataset.baizeai.io_datasets.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ spec:
* Note: syncMode can be "sync" (default) or "copy". "sync" removes files in destination that don't exist in source, "copy" only adds/updates files without removing existing ones.
- DATABASE: type(currently only support MySQL, other database types may be supported in the future.), host, port, dbName, tables(in the dbName), exportFormat(currently only support csv)
- HADOOP: coreSiteXml and hdfsSiteXml, sourcePath, username
- MANUAL:
type: object
x-kubernetes-preserve-unknown-fields: true
type:
Expand All @@ -243,6 +244,7 @@ spec:
- MODEL_SCOPE
- DATABASE
- HADOOP
- MANUAL
type: string
x-kubernetes-validations:
- message: Value is immutable
Expand All @@ -262,6 +264,7 @@ spec:
- MODEL_SCOPE: modelscope://<namespace>/<model>
- DATABASE: database://<ip>:<port>
- HADOOP: hdfs://<ip>:<port>
- MANUAL: manual://
type: string
x-kubernetes-validations:
- message: Value is immutable
Expand Down
39 changes: 32 additions & 7 deletions internal/controller/dataset/dataset_controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ func (r *DatasetReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
return ctrl.Result{}, nil
}

status := ds.Status.DeepCopy()
prevStatus := ds.Status.DeepCopy()
var reconcilers []reconciler
if kubeutils.IsDeleted(ds) {
reconcilers = []reconciler{
Expand Down Expand Up @@ -144,8 +144,8 @@ func (r *DatasetReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
}
resOk := ctrl.Result{}

if !reflect.DeepEqual(ds.Status, *status) {
err := r.Status().Update(ctx, ds)
if !reflect.DeepEqual(ds.Status, *prevStatus) {
err := r.updateStatus(ctx, ds, prevStatus)
if err != nil {
log.Errorf("error update status for %s/%s: %v", ds.Namespace, ds.Name, err)
return res30sec, err
Expand All @@ -162,6 +162,15 @@ func (r *DatasetReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
}
}

func (r *DatasetReconciler) updateStatus(ctx context.Context, ds *datasetv1alpha1.Dataset, prevStatus *datasetv1alpha1.DatasetStatus) error {
// Use a status-only merge patch instead of Status().Update. The dataset
// may have been updated while reconciling (for example, when adding the
// finalizer), and a full update would then fail on a stale resourceVersion.
statusBase := ds.DeepCopy()
statusBase.Status = *prevStatus
return r.Status().Patch(ctx, ds, client.MergeFrom(statusBase))
}

func supportPreload(ds *datasetv1alpha1.Dataset) bool {
switch ds.Spec.Source.Type {
case datasetv1alpha1.DatasetTypeGit,
Expand Down Expand Up @@ -544,8 +553,8 @@ func (r *DatasetReconciler) reconcileConfigMap(ctx context.Context, ds *datasetv

func (r *DatasetReconciler) reconcileJob(ctx context.Context, ds *datasetv1alpha1.Dataset) error {
if !supportPreload(ds) {
log.Infof("the type of %s/%s is %s not support preload, quit reconciling job",
ds.Namespace, ds.Name, ds.Spec.Source.Type)
// log.Infof("the type of %s/%s is %s not support preload, reconciling job is skipped",
// ds.Namespace, ds.Name, ds.Spec.Source.Type)
return nil
}
if kubeutils.IsDeleted(ds) {
Expand Down Expand Up @@ -840,8 +849,8 @@ func changeDefinitionForHadoop(sourceType datasetv1alpha1.DatasetType, jobSpec b
func (r *DatasetReconciler) reconcileJobStatus(ctx context.Context, ds *datasetv1alpha1.Dataset) error {
if !supportPreload(ds) {
ds.Status.LastSyncTime = ds.CreationTimestamp
log.Infof("the type of %s/%s is %s not support preload, quit reconciling job",
ds.Namespace, ds.Name, ds.Spec.Source.Type)
// log.Infof("the type of %s/%s is %s not support preload, reconciling job status is skipped",
// ds.Namespace, ds.Name, ds.Spec.Source.Type)
return nil
}
if !ds.Status.InProcessing {
Expand Down Expand Up @@ -912,6 +921,18 @@ func (r *DatasetReconciler) reconcilePhase(_ context.Context, ds *datasetv1alpha
ds.Status.Phase = datasetv1alpha1.DatasetStatusPhaseFailed
return nil
}
case datasetv1alpha1.DatasetTypeManual:
if _, ok := lo.Find(ds.Status.Conditions, func(c metav1.Condition) bool {
return c.Status == metav1.ConditionFalse
}); ok {
phase = datasetv1alpha1.DatasetStatusPhaseFailed
} else if ds.Status.PVCName == "" || !kubeutils.IsConditionReady(ds.Status.Conditions, condTypePVC) {
phase = datasetv1alpha1.DatasetStatusPhasePending
} else {
phase = datasetv1alpha1.DatasetStatusPhaseReady
}
ds.Status.Phase = phase
return nil
}

if ds.Spec.Source.Type == datasetv1alpha1.DatasetTypePVC {
Expand Down Expand Up @@ -943,6 +964,10 @@ func (r *DatasetReconciler) getSourceDataset(ctx context.Context, ds *datasetv1a
}

func (r *DatasetReconciler) validate(ctx context.Context, ds *datasetv1alpha1.Dataset) error {
if ds.Spec.Source.Type == datasetv1alpha1.DatasetTypeManual && ds.Spec.Source.URI != "manual://" {
return fmt.Errorf("MANUAL dataset source URI must be manual://")
}

if ds.Spec.Source.Type == datasetv1alpha1.DatasetTypeReference {
sourceDs, err := r.getSourceDataset(ctx, ds)
if err != nil {
Expand Down
189 changes: 189 additions & 0 deletions internal/controller/dataset/dataset_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -534,3 +534,192 @@ func TestDatasetReconciler_reconcilePVCNFSVersionDoesNotUpdateExistingPV(t *test
require.NoError(t, fakeClient.Get(context.Background(), client.ObjectKey{Name: pv.Name}, storedPV))
assert.Equal(t, []string{"nfsvers=4.1"}, storedPV.Spec.MountOptions)
}

func TestDatasetReconciler_reconcilePVCManual(t *testing.T) {
scheme := runtime.NewScheme()
require.NoError(t, datasetv1alpha1.AddToScheme(scheme))
require.NoError(t, corev1.AddToScheme(scheme))

storageClassName := "manual-storage"
ds := &datasetv1alpha1.Dataset{
ObjectMeta: metav1.ObjectMeta{
Name: "manual-dataset",
Namespace: "default",
UID: types.UID("manual-dataset-uid"),
},
Spec: datasetv1alpha1.DatasetSpec{
Source: datasetv1alpha1.DatasetSource{
Type: datasetv1alpha1.DatasetTypeManual,
URI: "manual://",
},
VolumeClaimTemplate: corev1.PersistentVolumeClaim{
Spec: corev1.PersistentVolumeClaimSpec{
AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
StorageClassName: &storageClassName,
Resources: corev1.VolumeResourceRequirements{
Requests: corev1.ResourceList{
corev1.ResourceStorage: resource.MustParse("1Gi"),
},
},
},
},
},
}
fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
reconciler := &DatasetReconciler{Client: fakeClient, Scheme: scheme}

require.NoError(t, reconciler.reconcilePVC(context.Background(), ds))
assert.Equal(t, ds.Name, ds.Status.PVCName)

pvc := &corev1.PersistentVolumeClaim{}
require.NoError(t, fakeClient.Get(context.Background(), client.ObjectKey{
Namespace: ds.Namespace,
Name: ds.Name,
}, pvc))
assert.Equal(t, []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce}, pvc.Spec.AccessModes)
require.NotNil(t, pvc.Spec.StorageClassName)
assert.Equal(t, storageClassName, *pvc.Spec.StorageClassName)
assert.True(t, pvc.Spec.Resources.Requests[corev1.ResourceStorage].Equal(resource.MustParse("1Gi")))
assert.Equal(t, ds.Name, pvc.Labels[constants.DatasetNameLabel])
require.Len(t, pvc.OwnerReferences, 1)
assert.Equal(t, ds.UID, pvc.OwnerReferences[0].UID)
}

func TestDatasetReconciler_updateStatusDoesNotRequireLatestMetadata(t *testing.T) {
scheme := runtime.NewScheme()
require.NoError(t, datasetv1alpha1.AddToScheme(scheme))

ctx := context.Background()
ds := &datasetv1alpha1.Dataset{
ObjectMeta: metav1.ObjectMeta{
Name: "manual-dataset",
Namespace: "public",
},
}
fakeClient := fake.NewClientBuilder().
WithScheme(scheme).
WithStatusSubresource(&datasetv1alpha1.Dataset{}).
WithObjects(ds).
Build()

stale := &datasetv1alpha1.Dataset{}
require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(ds), stale))

concurrent := &datasetv1alpha1.Dataset{}
require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(ds), concurrent))
concurrent.Labels = map[string]string{"updated": "concurrently"}
require.NoError(t, fakeClient.Update(ctx, concurrent))

previousStatus := stale.Status.DeepCopy()
stale.Status.Phase = datasetv1alpha1.DatasetStatusPhaseReady
reconciler := &DatasetReconciler{Client: fakeClient, Scheme: scheme}
require.NoError(t, reconciler.updateStatus(ctx, stale, previousStatus))

updated := &datasetv1alpha1.Dataset{}
require.NoError(t, fakeClient.Get(ctx, client.ObjectKeyFromObject(ds), updated))
assert.Equal(t, datasetv1alpha1.DatasetStatusPhaseReady, updated.Status.Phase)
assert.Equal(t, map[string]string{"updated": "concurrently"}, updated.Labels)
}

func TestManualDatasetDoesNotSupportPreload(t *testing.T) {
ds := &datasetv1alpha1.Dataset{
Spec: datasetv1alpha1.DatasetSpec{
Source: datasetv1alpha1.DatasetSource{
Type: datasetv1alpha1.DatasetTypeManual,
URI: "manual://",
},
DataSyncRound: 1,
},
}

assert.False(t, supportPreload(ds))
require.NoError(t, (&DatasetReconciler{}).reconcileJob(context.Background(), ds))
assert.False(t, ds.Status.InProcessing)
}

func TestDatasetReconciler_reconcilePhaseManual(t *testing.T) {
tests := []struct {
name string
status datasetv1alpha1.DatasetStatus
want datasetv1alpha1.DatasetStatusPhase
}{
{
name: "ready after PVC reconciliation succeeds",
status: datasetv1alpha1.DatasetStatus{
PVCName: "manual-dataset",
Conditions: []metav1.Condition{{
Type: condTypePVC,
Status: metav1.ConditionTrue,
}},
},
want: datasetv1alpha1.DatasetStatusPhaseReady,
},
{
name: "pending before PVC reconciliation succeeds",
status: datasetv1alpha1.DatasetStatus{
PVCName: "manual-dataset",
},
want: datasetv1alpha1.DatasetStatusPhasePending,
},
{
name: "failed when reconciliation reports an error",
status: datasetv1alpha1.DatasetStatus{
PVCName: "manual-dataset",
Conditions: []metav1.Condition{{
Type: condTypePVC,
Status: metav1.ConditionFalse,
}},
},
want: datasetv1alpha1.DatasetStatusPhaseFailed,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ds := &datasetv1alpha1.Dataset{
Spec: datasetv1alpha1.DatasetSpec{
Source: datasetv1alpha1.DatasetSource{
Type: datasetv1alpha1.DatasetTypeManual,
URI: "manual://",
},
DataSyncRound: 1,
},
Status: tt.status,
}

require.NoError(t, (&DatasetReconciler{}).reconcilePhase(context.Background(), ds))
assert.Equal(t, tt.want, ds.Status.Phase)
})
}
}

func TestDatasetReconciler_validateManualURI(t *testing.T) {
tests := []struct {
name string
uri string
wantErr bool
}{
{name: "canonical URI", uri: "manual://"},
{name: "non-canonical URI", uri: "manual://invalid", wantErr: true},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ds := &datasetv1alpha1.Dataset{
Spec: datasetv1alpha1.DatasetSpec{
Source: datasetv1alpha1.DatasetSource{
Type: datasetv1alpha1.DatasetTypeManual,
URI: tt.uri,
},
},
}

err := (&DatasetReconciler{}).validate(context.Background(), ds)
if tt.wantErr {
require.EqualError(t, err, "MANUAL dataset source URI must be manual://")
} else {
require.NoError(t, err)
}
})
}
}