Support v1beta2 operator in KubeflowTrainingService (#1058)

src/nni_manager/config/kubeflow/pytorchjob-crd-v1beta2.json

+{
+  "kind": "CustomResourceDefinition", 
+  "spec": {
+    "scope": "Namespaced", 
+    "version": "v1beta2", 
+    "group": "kubeflow.org", 
+    "names": {
+      "kind": "PyTorchJob", 
+      "plural": "pytorchjobs", 
+      "singular": "pytorchjob"
+    }
+  }, 
+  "apiVersion": "apiextensions.k8s.io/v1beta2", 
+  "metadata": {
+    "name": "pytorchjobs.kubeflow.org"
+  }
+}

src/nni_manager/config/kubeflow/tfjob-crd-v1beta2.json

+{
+  "kind": "CustomResourceDefinition", 
+  "spec": {
+    "scope": "Namespaced", 
+    "version": "v1beta2", 
+    "group": "kubeflow.org", 
+    "names": {
+      "kind": "TFJob", 
+      "plural": "tfjobs", 
+      "singular": "tfjob"
+    }
+  }, 
+  "apiVersion": "apiextensions.k8s.io/v1beta2", 
+  "metadata": {
+    "name": "tfjobs.kubeflow.org"
+  }
+}

src/nni_manager/training_service/kubernetes/kubeflow/kubeflowApiClient.ts

@@ -29,20 +29,35 @@ abstract class KubeflowOperatorClient extends KubernetesCRDClient{
      */
     public static generateOperatorClient(kubeflowOperator: KubeflowOperator, 
                                     operatorApiVersion: string): KubernetesCRDClient {
-        if(kubeflowOperator === 'tf-operator') {
-            if(operatorApiVersion == 'v1alpha2') {
-                return new TFOperatorClientV1Alpha2();
-            } else if(operatorApiVersion == 'v1beta1') {
-                return new TFOperatorClientV1Beta1();
-            }
-        } else if(kubeflowOperator === 'pytorch-operator') {
-            if(operatorApiVersion == 'v1alpha2') {
-                return new PytorchOperatorClientV1Alpha2();
-            } else if(operatorApiVersion == 'v1beta1') {
-                return new PytorchOperatorClientV1Beta1();
+        switch(kubeflowOperator) {
+            case 'tf-operator': {
+                switch(operatorApiVersion) {
+                    case 'v1alpha2': {
+                        return new TFOperatorClientV1Alpha2();
+                    }
+                    case 'v1beta1': {
+                        return new TFOperatorClientV1Beta1();
+                    }
+                    case 'v1beta2': {
+                        return new TFOperatorClientV1Beta2();
+                    }
+                }
+                break;
             }
+            case 'pytorch-operator': {
+                switch(operatorApiVersion) {
+                    case 'v1alpha2': {
+                        return new PyTorchOperatorClientV1Alpha2();
+                    }
+                    case 'v1beta1': {
+                        return new PyTorchOperatorClientV1Beta1();
+                    }
+                    case 'v1beta2': {
+                        return new PyTorchOperatorClientV1Beta2();
+                    }
+                }
+            } 
         }
-
         throw new Error(`Invalid operator ${kubeflowOperator} or apiVersion ${operatorApiVersion}`);
     }
 }
@@ -85,7 +100,26 @@ class TFOperatorClientV1Beta1 extends KubernetesCRDClient {
     }    
 }
 
-class PytorchOperatorClientV1Alpha2 extends KubeflowOperatorClient {
+class TFOperatorClientV1Beta2 extends KubernetesCRDClient {
+    /**
+     * constructor, to initialize tfjob CRD definition
+     */
+    public constructor() {
+        super();
+        this.crdSchema = JSON.parse(fs.readFileSync('./config/kubeflow/tfjob-crd-v1beta2.json', 'utf8'));
+        this.client.addCustomResourceDefinition(this.crdSchema);
+    }
+
+    protected get operator(): any {
+        return this.client.apis["kubeflow.org"].v1beta2.namespaces('default').tfjobs;
+    }
+
+    public get containerName(): string {
+        return 'tensorflow';
+    }    
+}
+
+class PyTorchOperatorClientV1Alpha2 extends KubeflowOperatorClient {
     /**
      * constructor, to initialize tfjob CRD definition
      */
@@ -104,7 +138,7 @@ class PytorchOperatorClientV1Alpha2 extends KubeflowOperatorClient {
     }
 }
 
-class PytorchOperatorClientV1Beta1 extends KubernetesCRDClient {
+class PyTorchOperatorClientV1Beta1 extends KubernetesCRDClient {
     /**
      * constructor, to initialize tfjob CRD definition
      */
@@ -123,5 +157,24 @@ class PytorchOperatorClientV1Beta1 extends KubernetesCRDClient {
     }
 }
 
+class PyTorchOperatorClientV1Beta2 extends KubernetesCRDClient {
+    /**
+     * constructor, to initialize tfjob CRD definition
+     */
+    public constructor() {
+        super();
+        this.crdSchema = JSON.parse(fs.readFileSync('./config/kubeflow/pytorchjob-crd-v1beta2.json', 'utf8'));
+        this.client.addCustomResourceDefinition(this.crdSchema);
+    }
+
+    protected get operator(): any {
+        return this.client.apis["kubeflow.org"].v1beta2.namespaces('default').pytorchjobs;
+    }
+
+    public get containerName(): string {
+        return 'pytorch';
+    }
+}
+
 export { KubeflowOperatorClient, GeneralK8sClient };
 

src/nni_manager/training_service/kubernetes/kubeflow/kubeflowConfig.ts

@@ -28,7 +28,7 @@ import { MethodNotImplementedError } from '../../../common/errors';
 export type KubeflowOperator = 'tf-operator' | 'pytorch-operator' ;
 export type DistTrainRole = 'worker' | 'ps' | 'master';
 export type KubeflowJobStatus = 'Created' | 'Running' | 'Failed' | 'Succeeded';
-export type OperatorApiVersion = 'v1alpha2' | 'v1beta1';
+export type OperatorApiVersion = 'v1alpha2' | 'v1beta1' | 'v1beta2';
 
 export class KubeflowClusterConfig extends KubernetesClusterConfig {
     public readonly operator: KubeflowOperator;