fix(sdk): Fix type annotation for train method's trainer parameter (#2646)

* SDK: Fix type annotation of `train` method's `trainer` parameter

Signed-off-by: Eoin Fennessy <efenness@redhat.com>

* Use `Optional[Union[...]] = None`

Signed-off-by: Eoin Fennessy <efenness@redhat.com>

---------

Signed-off-by: Eoin Fennessy <efenness@redhat.com>

Author: eoinfennessy

sdk/kubeflow/trainer/api/trainer_client.py

@@ -18,7 +18,7 @@
 import random
 import string
 import uuid
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import kubeflow.trainer.models as models
 from kubeflow.trainer.constants import constants
@@ -153,20 +153,21 @@ def train(
         self,
         runtime: types.Runtime = types.DEFAULT_RUNTIME,
         initializer: Optional[types.Initializer] = None,
-        trainer: Optional[types.CustomTrainer] = None,
+        trainer: Optional[Union[types.CustomTrainer, types.BuiltinTrainer]] = None,
     ) -> str:
         """
         Create the TrainJob. You can configure these types of training task:
 
         - Custom Training Task: Training with a self-contained function that encapsulates
             the entire model training process, e.g. `CustomTrainer`.
+        - Builtin Training Task: Configures a post-training job using torchtune.
 
         Args:
             runtime (`types.Runtime`): Reference to one of existing Runtimes.
             initializer (`Optional[types.Initializer]`):
                 Configuration for the dataset and model initializers.
-            trainer (`Optional[types.CustomTrainer]`):
-                Configuration for Custom Training Task.
+            trainer (`Union[types.CustomTrainer, types.BuiltinTrainer, None]`):
+                Configuration for Custom or Builtin Training Task.
 
         Returns:
             str: The unique name of the TrainJob that has been generated.