From a449ecec6924e5024a58e8fd82938ebd0caf9deb Mon Sep 17 00:00:00 2001
From: deepanker13 <deepanker.gupta@nutanix.com>
Date: Wed, 10 Jan 2024 15:37:27 +0530
Subject: [PATCH] adding extra requires in setup.py, fixinf ci

---
 .github/workflows/integration-tests.yaml      |  1 -
 .github/workflows/test-python.yaml            |  4 +--
 sdk/python/Dockerfile.conformance             |  3 +-
 .../kubeflow/training/api/training_client.py  | 32 +++++++++++++------
 sdk/python/setup.py                           |  5 ++-
 5 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml
index f9cec9e723..4eb63428bc 100644
--- a/.github/workflows/integration-tests.yaml
+++ b/.github/workflows/integration-tests.yaml
@@ -109,7 +109,6 @@ jobs:
       - name: Run tests
         run: |
           pip install pytest
-          pip install -r sdk/python/kubeflow/storage_init_container/requirements.txt
           python3 -m pip install -e sdk/python; pytest sdk/python/test --log-cli-level=info --namespace=default
         env:
           GANG_SCHEDULER_NAME: ${{ matrix.gang-scheduler-name }}
diff --git a/.github/workflows/test-python.yaml b/.github/workflows/test-python.yaml
index 2e18075c1e..28994f2d07 100644
--- a/.github/workflows/test-python.yaml
+++ b/.github/workflows/test-python.yaml
@@ -25,7 +25,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install pytest python-dateutil urllib3 kubernetes
-          pip install -r ./sdk/python/kubeflow/storage_init_container/requirements.txt
-
+          pip install -U './sdk/python[huggingface]'
+          
       - name: Run unit test for training sdk
         run: pytest ./sdk/python/kubeflow/training/api/training_client_test.py
\ No newline at end of file
diff --git a/sdk/python/Dockerfile.conformance b/sdk/python/Dockerfile.conformance
index b3a1e2d5f3..4cdc6f3bd6 100644
--- a/sdk/python/Dockerfile.conformance
+++ b/sdk/python/Dockerfile.conformance
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Dockerfile for building the source code of conformance tests
-FROM python:3.10
+FROM python:3.10-alpine
 
 WORKDIR /kubeflow/training-operator/sdk/python
 
@@ -25,7 +25,6 @@ COPY sdk/python/conformance/run.sh .
 RUN chmod +x run.sh
 
 RUN pip install pytest
-RUN pip install -r ./kubeflow/storage_init_container/requirements.txt
 RUN python -m pip install -e .
 
 ENTRYPOINT [ "./run.sh" ]
\ No newline at end of file
diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
index 0cfa0d9072..5f17e641dd 100644
--- a/sdk/python/kubeflow/training/api/training_client.py
+++ b/sdk/python/kubeflow/training/api/training_client.py
@@ -25,13 +25,8 @@
 from kubeflow.training.constants import constants
 from kubeflow.training.utils import utils
 
-from kubeflow.storage_init_container.s3 import S3DatasetParams
-from kubeflow.storage_init_container.hugging_face import (
-    HuggingFaceModelParams,
-    HuggingFaceTrainParams,
-    HfDatasetParams,
-    INIT_CONTAINER_MOUNT_PATH,
-)
+from kubeflow.storage_init_container.constants import INIT_CONTAINER_MOUNT_PATH
+
 
 logger = logging.getLogger(__name__)
 
@@ -103,15 +98,32 @@ def train(
         num_workers: int = 1,
         num_procs_per_worker: int = 1,
         storage_config: Dict[str, str] = {"size": "10Gi", "storage_class": None},
-        model_provider_parameters: HuggingFaceModelParams = None,
-        dataset_provider_parameters: Union[HfDatasetParams, S3DatasetParams] = None,
-        train_parameters: HuggingFaceTrainParams = None,
+        model_provider_parameters=None,
+        dataset_provider_parameters=None,
+        train_parameters=None,
         resources_per_worker: Union[dict, client.V1ResourceRequirements, None] = None,
         # Dict[Literal["gpu", "cpu", "memory"], any] = None,
     ):
         """
         Higher level train api
+        model_provider_parameters: It can be of type HuggingFaceModelParams
+        dataset_provider_parameters: It can be of type HfDatasetParams or S3DatasetParams
+        train_parameters: It can be of type HuggingFaceTrainParams
         """
+        try:
+            import peft
+            import transformers
+        except ImportError:
+            print(
+                "train api dependencies not installed. Run pip install -U 'kubeflow-training[huggingface]' "
+            )
+        from kubeflow.storage_init_container.s3 import S3DatasetParams
+        from kubeflow.storage_init_container.hugging_face import (
+            HuggingFaceModelParams,
+            HuggingFaceTrainParams,
+            HfDatasetParams,
+        )
+
         if (
             not name
             or not model_provider_parameters
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
index c584fa38ce..c71d90919a 100644
--- a/sdk/python/setup.py
+++ b/sdk/python/setup.py
@@ -64,5 +64,8 @@
     ],
     install_requires=REQUIRES,
     tests_require=TESTS_REQUIRES,
-    extras_require={"test": TESTS_REQUIRES},
+    extras_require={
+        "test": TESTS_REQUIRES,
+        "huggingface": ["transformers>=4.20.0", "peft>=0.3.0"],
+    },
 )