From 0b86d0000eb854865554f2b987f34f9cef165cf9 Mon Sep 17 00:00:00 2001
From: machichima <60069744+machichima@users.noreply.github.com>
Date: Sat, 14 Dec 2024 10:53:36 +0800
Subject: [PATCH] Example using ImageSpec in Raw container (#1773)

Signed-off-by: machichima <nary12321@gmail.com>
---
 .../calculate-ellipse-area-new.py             | 29 +++++++++++++++++++
 .../customizing_dependencies/raw_container.py | 12 ++++++--
 2 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 examples/customizing_dependencies/customizing_dependencies/calculate-ellipse-area-new.py

diff --git a/examples/customizing_dependencies/customizing_dependencies/calculate-ellipse-area-new.py b/examples/customizing_dependencies/customizing_dependencies/calculate-ellipse-area-new.py
new file mode 100644
index 000000000..7c589da7c
--- /dev/null
+++ b/examples/customizing_dependencies/customizing_dependencies/calculate-ellipse-area-new.py
@@ -0,0 +1,29 @@
+import math
+import sys
+
+
+def write_output(output_dir, output_file, v):
+    with open(f"{output_dir}/{output_file}", "w") as f:
+        f.write(str(v))
+
+
+def calculate_area(a, b):
+    return math.pi * a * b
+
+
+def main(a, b, output_dir):
+    a = float(a)
+    b = float(b)
+
+    area = calculate_area(a, b)
+
+    write_output(output_dir, "area", area)
+    write_output(output_dir, "metadata", "[from python rawcontainer]")
+
+
+if __name__ == "__main__":
+    a = sys.argv[1]
+    b = sys.argv[2]
+    output_dir = sys.argv[3]
+
+    main(a, b, output_dir)
diff --git a/examples/customizing_dependencies/customizing_dependencies/raw_container.py b/examples/customizing_dependencies/customizing_dependencies/raw_container.py
index 9d986da56..6e579283f 100644
--- a/examples/customizing_dependencies/customizing_dependencies/raw_container.py
+++ b/examples/customizing_dependencies/customizing_dependencies/raw_container.py
@@ -1,6 +1,6 @@
 import logging
 
-from flytekit import ContainerTask, kwtypes, task, workflow
+from flytekit import ContainerTask, ImageSpec, kwtypes, task, workflow
 from flytekit.core.base_task import TaskMetadata
 
 logger = logging.getLogger(__file__)
@@ -29,16 +29,22 @@
     metadata=TaskMetadata(cache=True, cache_version="1.0"),
 )
 
+# use `ImageSpec` to copy files or directories into container `/root`
 calculate_ellipse_area_python = ContainerTask(
     name="ellipse-area-metadata-python",
     input_data_dir="/var/inputs",
     output_data_dir="/var/outputs",
     inputs=kwtypes(a=float, b=float),
     outputs=kwtypes(area=float, metadata=str),
-    image="ghcr.io/flyteorg/rawcontainers-python:v2",
+    image=ImageSpec(
+        base_image="ghcr.io/flyteorg/rawcontainers-python:v2",
+        registry="localhost:30000",
+        builder="default",
+        copy=["calculate-ellipse-area-new.py"],
+    ),
     command=[
         "python",
-        "calculate-ellipse-area.py",
+        "calculate-ellipse-area-new.py",
         "{{.inputs.a}}",
         "{{.inputs.b}}",
         "/var/outputs",