feat(envd): Add switch for custom envd context (flyteorg#2500)

Signed-off-by: Kevin Su <pingsutw@apache.org>
MortalHappiness · Jun 20, 2024 · b98e941 · b98e941
1 parent 32bab30
commit b98e941
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py
@@ -14,6 +14,7 @@
 from flytekit.tools.ignore import DockerIgnore, GitIgnore, IgnoreGroup, StandardIgnore
 
 FLYTE_LOCAL_REGISTRY = "localhost:30000"
+FLYTE_ENVD_CONTEXT = "FLYTE_ENVD_CONTEXT"
 
 
 class EnvdImageSpecBuilder(ImageSpecBuilder):
@@ -38,6 +39,9 @@ def build_image(self, image_spec: ImageSpec):
 
 
 def envd_context_switch(registry: str):
+    if os.getenv(FLYTE_ENVD_CONTEXT):
+        execute_command(f"envd context use --name {os.getenv(FLYTE_ENVD_CONTEXT)}")
+        return
     if registry == FLYTE_LOCAL_REGISTRY:
         # Assume buildkit daemon is running within the sandbox and exposed on port 30003
         command = "envd context create --name flyte-sandbox --builder tcp --builder-address localhost:30003 --use"