template bucket info for MNIST dataset

Signed-off-by: Kevin <[email protected]>

Author: KPostOffice

tests/kfto/resources/kfto_sdk_mnist.py

@@ -6,6 +6,10 @@ def train_func():
     from torchvision import datasets, transforms
     import torch.distributed as dist
     from pathlib import Path
+    from minio import Minio
+    import shutil
+    import gzip
+
 
     # [1] Setup PyTorch DDP. Distributed environment will be set automatically by Training Operator.
     dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")
@@ -45,13 +49,63 @@ def forward(self, x):
     optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
 
     # [4] Setup FashionMNIST dataloader and distribute data across PyTorchJob workers.
-    Path(f"./data{local_rank}").mkdir(exist_ok=True)
-    dataset = datasets.FashionMNIST(
-        f"./data{local_rank}",
-        download=True,
-        train=True,
-        transform=transforms.Compose([transforms.ToTensor()]),
-    )
+    dataset_path = "./data"
+    dataset_dir = os.path.join(dataset_path, "MNIST/raw")
+    with_aws = "{{.StorageBucketNameExists}}"
+    endpoint = "{{.StorageBucketDefaultEndpoint}}"
+    access_key = "{{.StorageBucketAccessKeyId}}"
+    secret_key = "{{.StorageBucketSecretKey}}"
+    bucket_name = "{{.StorageBucketName}}"
+    prefix = "{{.StorageBucketMnistDir}}"
+    if with_aws != "true":
+        client = Minio(
+            endpoint,
+            access_key=access_key,
+            secret_key=secret_key,
+            cert_check=False,
+            secure=False,  #TODO
+        )
+
+        if not os.path.exists(dataset_dir):
+            os.makedirs(dataset_dir)
+
+        for item in client.list_objects(
+            bucket_name, prefix=prefix, recursive=True
+        ):  
+            file_name=item.object_name[len(prefix)+1:]
+            dataset_file_path = os.path.join(dataset_dir, file_name)
+            print(f"Downloading dataset file {file_name} to {dataset_file_path}..")
+            if not os.path.exists(dataset_file_path):
+                client.fget_object(
+                    bucket_name, item.object_name, dataset_file_path
+                )
+                # Unzip files -- 
+                ## Sample zipfilepath : ../data/MNIST/raw/t10k-images-idx3-ubyte.gz
+                with gzip.open(dataset_file_path, "rb") as f_in:
+                    filename=file_name.split(".")[0]    #-> t10k-images-idx3-ubyte
+                    file_path=("/".join(dataset_file_path.split("/")[:-1]))     #->../data/MNIST/raw
+                    full_file_path=os.path.join(file_path,filename)     #->../data/MNIST/raw/t10k-images-idx3-ubyte
+                    print(f"Extracting {dataset_file_path} to {file_path}..")
+
+                    with open(full_file_path, "wb") as f_out:
+                        shutil.copyfileobj(f_in, f_out)
+                    print(f"Dataset file downloaded : {full_file_path}\n")
+                # delete zip file
+                os.remove(dataset_file_path)
+
+        dataset = datasets.MNIST(
+            dataset_path,
+            train=True,
+            download=False,
+            transform=transforms.Compose([transforms.ToTensor()]),
+        )
+    else:
+        dataset = datasets.MNIST(
+            dataset_path,
+            train=True,
+            download=True,
+            transform=transforms.Compose([transforms.ToTensor()]),
+        )
     train_loader = torch.utils.data.DataLoader(
         dataset=dataset,
         batch_size=128,

tests/kfto/resources/mnist_kfto.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "b55bc3ea-4ce3-49bf-bb1f-e209de8ca47a",
    "metadata": {
     "tags": []
@@ -19,7 +19,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 7,
    "id": "72dd1751",
    "metadata": {},
    "outputs": [],
@@ -33,14 +33,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "4ca70b20",
    "metadata": {},
    "outputs": [],
    "source": [
-    "config = c.Configuration(host=openshift_api_url, api_key=token)\n",
-    "config.verify_ssl = False\n",
-    "tc = TrainingClient()"
+    "# config = c.Configuration(host=openshift_api_url, api_key=token)\n",
+    "# config.verify_ssl = False\n",
+    "tc = TrainingClient()\n",
+    "# config.api_key_prefix"
    ]
   },
   {
@@ -50,17 +51,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "try:\n",
-    "   tc.create_job(\n",
-    "      name=\"pytorch-ddp\",\n",
-    "      namespace=namespace,\n",
-    "      train_func=train_func,\n",
-    "      num_workers=2,\n",
-    "      resources_per_worker={\"gpu\": num_gpus},\n",
-    "      base_image=\"quay.io/kpostlet/torch-train:withvision\",\n",
-    "   )\n",
-    "except:\n",
-    "   pass"
+    "tc.create_job(\n",
+    "   name=\"pytorch-ddp\",\n",
+    "   namespace=namespace,\n",
+    "   train_func=train_func,\n",
+    "   num_workers=2,\n",
+    "   resources_per_worker={\"gpu\": num_gpus},\n",
+    "   base_image=\"quay.io/kpostlet/torch-train:with-minivision\",\n",
+    "   # packages_to_install=[\"torchvision==0.19.0\", \"--target=/tmp/lib\"],\n",
+    "   # env_vars={\"PYTHONPATH\": \"/tmp/lib:$PYTHONPATH\", \"NCCL_DEBUG\": \"INFO\", \"TORCH_DISTRIBUTED_DEBUG\": \"DETAIL\"}\n",
+    ")"
    ]
   },
   {