Add Python SDK for Kubeflow Training Operator

README.md

@@ -47,6 +47,15 @@ For users who prefer to use original tensorflow controllers, please checkout v1.
 kubectl apply -k "github.com/kubeflow/tf-operator.git/manifests/overlays/standalone?ref=v1.2.0"
 ```
 
+### Python SDK for Kubeflow Training Operator
+
+Training Operator provides Python SDK for the custom resources. More docs are available in [sdk/python](sdk/python) folder.
+
+Use `pip install` command to install the latest release of the SDK:
+```
+pip install kubeflow-training
+```
+
 ## Quick Start
 
 Please refer to the [quick-start-v1.md](docs/quick-start-v1.md) and [Kubeflow Training User Guide](https://www.kubeflow.org/docs/guides/components/tftraining/) for more information.

docs/development/developer_guide.md

@@ -80,6 +80,21 @@ kubectl create -f ./tf_job_mnist.yaml
 
 On ubuntu the default go package appears to be gccgo-go which has problems see [issue](https://github.com/golang/go/issues/15429) golang-go package is also really old so install from golang tarballs instead.
 
+## Generate Python SDK
+
+To generate Python SDK for the operator, run:
+```
+.hack/python-sdk/gen-sdk.sh
+```
+This command will re-generate the api and model files together with the documentation and model tests.
+The following files/folders in `sdk/python` are auto-generated and should not be modified directly:
+```
+docs
+kubeflow/training/models
+kubeflow/training/*.py
+test/*.py
+```
+
 ## Code Style
 
 ### Python

hack/python-sdk/gen-sdk.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 
-# Copyright 2019 The Kubeflow Authors.
+# Copyright 2021 The Kubeflow Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,10 +18,12 @@ set -o errexit
 set -o nounset
 set -o pipefail
 
+repo_root="$(realpath "$(dirname "$0")/../..")"
+
 SWAGGER_JAR_URL="https://repo1.maven.org/maven2/org/openapitools/openapi-generator-cli/4.3.1/openapi-generator-cli-4.3.1.jar"
-SWAGGER_CODEGEN_JAR="hack/python-sdk/openapi-generator-cli.jar"
-SWAGGER_CODEGEN_CONF="hack/python-sdk/swagger_config.json"
-SDK_OUTPUT_PATH="/tmp/sdk/python"
+SWAGGER_CODEGEN_JAR="${repo_root}/hack/python-sdk/openapi-generator-cli.jar"
+SWAGGER_CODEGEN_CONF="${repo_root}/hack/python-sdk/swagger_config.json"
+SDK_OUTPUT_PATH="${repo_root}/sdk/python"
 FRAMEWORKS=(tensorflow pytorch mxnet xgboost)
 VERSION=1.3.0
 
@@ -32,14 +34,16 @@ fi
 echo "Generating OpenAPI specification ..."
 echo "./hack/update-codegen.sh already help us generate openapi specs ..."
 
-echo "Downloading the swagger-codegen JAR package ..."
-wget -O ${SWAGGER_CODEGEN_JAR} ${SWAGGER_JAR_URL}
+if [[ ! -f "$SWAGGER_CODEGEN_JAR" ]]; then
+  echo "Downloading the swagger-codegen JAR package ..."
+  wget -O "${SWAGGER_CODEGEN_JAR}" ${SWAGGER_JAR_URL}
+fi
 
 
 for FRAMEWORK in ${FRAMEWORKS[@]}; do
     SWAGGER_CODEGEN_FILE="pkg/apis/${FRAMEWORK}/v1/swagger.json"
     echo "Generating swagger file for ${FRAMEWORK} ..."
-    go run hack/python-sdk/main.go ${FRAMEWORK} ${VERSION} > ${SWAGGER_CODEGEN_FILE}
+    go run "${repo_root}"/hack/python-sdk/main.go "${FRAMEWORK}" ${VERSION} > "${SWAGGER_CODEGEN_FILE}"
 done
 
 echo "Merging swagger files from different frameworks into one"
@@ -50,11 +54,17 @@ chmod +x /tmp/swagger
 
 # it will report warning like 'v1.SchedulingPolicy' already exists in primary or higher priority mixin, skipping
 # error code is not 0 but t's acceptable.
-/tmp/swagger mixin pkg/apis/tensorflow/v1/swagger.json pkg/apis/pytorch/v1/swagger.json pkg/apis/mxnet/v1/swagger.json pkg/apis/xgboost/v1/swagger.json \
---output hack/python-sdk/swagger.json --quiet || true
+/tmp/swagger mixin "${repo_root}"/pkg/apis/tensorflow/v1/swagger.json "${repo_root}"/pkg/apis/pytorch/v1/swagger.json \
+  "${repo_root}"/pkg/apis/mxnet/v1/swagger.json "${repo_root}"/pkg/apis/xgboost/v1/swagger.json \
+  --output "${repo_root}"/hack/python-sdk/swagger.json --quiet || true
 
-echo "Generating Python SDK for ${FRAMEWORK} ..."
-java -jar ${SWAGGER_CODEGEN_JAR} generate -i hack/python-sdk/swagger.json -g python -o ${SDK_OUTPUT_PATH} -c ${SWAGGER_CODEGEN_CONF}
+echo "Removing previously generated files ..."
+rm -rf "${SDK_OUTPUT_PATH}"/docs "${SDK_OUTPUT_PATH}"/kubeflow/training/models "${SDK_OUTPUT_PATH}"/kubeflow/training/*.py "${SDK_OUTPUT_PATH}"/test/*.py
+echo "Generating Python SDK for Training Operator ..."
+java -jar "${SWAGGER_CODEGEN_JAR}" generate -i "${repo_root}"/hack/python-sdk/swagger.json -g python -o "${SDK_OUTPUT_PATH}" -c "${SWAGGER_CODEGEN_CONF}"
 
 echo "Kubeflow Training Operator Python SDK is generated successfully to folder ${SDK_OUTPUT_PATH}/."
-rm /tmp/swagger
+rm /tmp/swagger
+
+echo "Running post-generation script ..."
+"${repo_root}"/hack/python-sdk/post_gen.py

hack/python-sdk/main.go

@@ -90,6 +90,9 @@ func main() {
 func swaggify(name, framework string) string {
 	name = strings.Replace(name, fmt.Sprintf("github.com/kubeflow/tf-operator/pkg/apis/%s/", framework), "", -1)
 	name = strings.Replace(name, "github.com/kubeflow/common/pkg/apis/common/", "", -1)
+	name = strings.Replace(name, "k8s.io/api/core/", "", -1)
+	name = strings.Replace(name, "k8s.io/apimachinery/pkg/apis/meta/", "", -1)
+	name = strings.Replace(name, "k8s.io/apimachinery/pkg/api/resource", "", -1)
 	name = strings.Replace(name, "/", ".", -1)
 	return name
 }

hack/python-sdk/post_gen.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+"""
+This script is used for updating generated SDK files.
+"""
+
+import os
+import fileinput
+import re
+
+__replacements = [
+    ("import kubeflow.training", "from kubeflow.training.models import *"),
+    ("kubeflow.training.models.v1\/.*.v1.", "V1")
+]
+
+sdk_dir = os.path.abspath(os.path.join(__file__, "../../..", "sdk/python"))
+
+
+def main():
+    fix_test_files()
+
+
+def fix_test_files() -> None:
+    """
+    Fix invalid model imports in generated model tests
+    """
+    os.path.realpath(__file__)
+    test_folder_dir = os.path.join(sdk_dir, "test")
+    test_files = os.listdir(test_folder_dir)
+    for test_file in test_files:
+        print(f"Precessing file {test_file}")
+        if test_file.endswith(".py"):
+            with fileinput.FileInput(os.path.join(test_folder_dir, test_file), inplace=True) as file:
+                for line in file:
+                    print(_apply_regex(line), end='')
+
+
+def _apply_regex(input_str: str) -> str:
+    for pattern, replacement in __replacements:
+        input_str = re.sub(pattern, replacement, input_str)
+    return input_str
+
+
+if __name__ == '__main__':
+    main()

hack/python-sdk/swagger.json

@@ -62,7 +62,7 @@
         },
         "lastReconcileTime": {
           "description": "Represents last time when the job was reconciled. It is not guaranteed to be set in happens-before order across separate operations. It is represented in RFC3339 form and is in UTC.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         },
         "replicaStatuses": {
           "description": "ReplicaStatuses is map of ReplicaType and ReplicaStatus, specifies the status of each replica.",
@@ -73,7 +73,7 @@
         },
         "startTime": {
           "description": "Represents time when the job was acknowledged by the job controller. It is not guaranteed to be set in happens-before order across separate operations. It is represented in RFC3339 form and is in UTC.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         }
       }
     },
@@ -90,7 +90,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ObjectMeta"
+          "$ref": "#/definitions/v1.ObjectMeta"
         },
         "spec": {
           "$ref": "#/definitions/v1.MXJobSpec"
@@ -122,7 +122,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ListMeta"
+          "$ref": "#/definitions/v1.ListMeta"
         }
       }
     },
@@ -169,7 +169,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ObjectMeta"
+          "$ref": "#/definitions/v1.ObjectMeta"
         },
         "spec": {
           "description": "Specification of the desired state of the PyTorchJob.",
@@ -205,7 +205,7 @@
         },
         "metadata": {
           "description": "Standard list metadata.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ListMeta"
+          "$ref": "#/definitions/v1.ListMeta"
         }
       }
     },
@@ -245,7 +245,7 @@
         },
         "template": {
           "description": "Template is the object that describes the pod that will be created for this replica. RestartPolicy in PodTemplateSpec will be overide by RestartPolicy in ReplicaSpec",
-          "$ref": "#/definitions/k8s.io.api.core.v1.PodTemplateSpec"
+          "$ref": "#/definitions/v1.PodTemplateSpec"
         }
       }
     },
@@ -310,7 +310,7 @@
         "minResources": {
           "type": "object",
           "additionalProperties": {
-            "$ref": "#/definitions/k8s.io.apimachinery.pkg.api.resource.Quantity"
+            "$ref": "#/definitions/.Quantity"
           }
         },
         "priorityClass": {
@@ -334,7 +334,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ObjectMeta"
+          "$ref": "#/definitions/v1.ObjectMeta"
         },
         "spec": {
           "description": "Specification of the desired state of the TFJob.",
@@ -473,4 +473,4 @@
       }
     }
   }
-}
+}

pkg/apis/mxnet/v1/swagger.json

@@ -17,11 +17,11 @@
       "properties": {
         "lastTransitionTime": {
           "description": "Last time the condition transitioned from one status to another.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         },
         "lastUpdateTime": {
           "description": "The last time this condition was updated.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         },
         "message": {
           "description": "A human readable message indicating details about the transition.",
@@ -51,7 +51,7 @@
       "properties": {
         "completionTime": {
           "description": "Represents time when the job was completed. It is not guaranteed to be set in happens-before order across separate operations. It is represented in RFC3339 form and is in UTC.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         },
         "conditions": {
           "description": "Conditions is an array of current observed job conditions.",
@@ -62,7 +62,7 @@
         },
         "lastReconcileTime": {
           "description": "Represents last time when the job was reconciled. It is not guaranteed to be set in happens-before order across separate operations. It is represented in RFC3339 form and is in UTC.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         },
         "replicaStatuses": {
           "description": "ReplicaStatuses is map of ReplicaType and ReplicaStatus, specifies the status of each replica.",
@@ -73,7 +73,7 @@
         },
         "startTime": {
           "description": "Represents time when the job was acknowledged by the job controller. It is not guaranteed to be set in happens-before order across separate operations. It is represented in RFC3339 form and is in UTC.",
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.Time"
+          "$ref": "#/definitions/v1.Time"
         }
       }
     },
@@ -90,7 +90,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ObjectMeta"
+          "$ref": "#/definitions/v1.ObjectMeta"
         },
         "spec": {
           "$ref": "#/definitions/v1.MXJobSpec"
@@ -122,7 +122,7 @@
           "type": "string"
         },
         "metadata": {
-          "$ref": "#/definitions/k8s.io.apimachinery.pkg.apis.meta.v1.ListMeta"
+          "$ref": "#/definitions/v1.ListMeta"
         }
       }
     },
@@ -171,7 +171,7 @@
         },
         "template": {
           "description": "Template is the object that describes the pod that will be created for this replica. RestartPolicy in PodTemplateSpec will be overide by RestartPolicy in ReplicaSpec",
-          "$ref": "#/definitions/k8s.io.api.core.v1.PodTemplateSpec"
+          "$ref": "#/definitions/v1.PodTemplateSpec"
         }
       }
     },
@@ -236,7 +236,7 @@
         "minResources": {
           "type": "object",
           "additionalProperties": {
-            "$ref": "#/definitions/k8s.io.apimachinery.pkg.api.resource.Quantity"
+            "$ref": "#/definitions/.Quantity"
           }
         },
         "priorityClass": {