truefoundry 0.11.1rc1__py3-none-any.whl → 0.11.3__py3-none-any.whl

truefoundry/__init__.py

@@ -12,6 +12,7 @@ from truefoundry_sdk import (
     ToolSchema,
     UserMessage,
 )
+from truefoundry_sdk.client import TrueFoundry
 
 from truefoundry._client import client
 from truefoundry.common.warnings import (
@@ -39,6 +40,7 @@ __all__ = [
     "render_prompt",
     "suppress_truefoundry_deprecation_warnings",
     "SystemMessage",
+    "TrueFoundry",
     "ToolCall",
     "ToolMessage",
     "ToolSchema",

truefoundry/common/constants.py

@@ -68,6 +68,8 @@ class TrueFoundrySdkEnv(BaseSettings):
     # TODO(gw): Use another image with more linient rate limits
     TFY_SPARK_BUILD_SPARK_IMAGE_REPO: str = "public.ecr.aws/bitnami/spark"
 
+    TFY_TASK_PYSPARK_BUILD_SPARK_IMAGE_REPO: str = "public.ecr.aws/bitnami/spark"
+
     # For local development, this enables futher configuration via _TFYServersConfig
     TFY_CLI_LOCAL_DEV_MODE: bool = False
 

truefoundry/common/utils.py

@@ -16,6 +16,7 @@ from truefoundry.common.constants import (
     API_SERVER_RELATIVE_PATH,
     ENV_VARS,
     MLFOUNDRY_SERVER_RELATIVE_PATH,
+    TFY_API_KEY_ENV_KEY,
     TFY_DEBUG_ENV_KEY,
     TFY_HOST_ENV_KEY,
     TFY_INTERNAL_ENV_KEY,
@@ -113,9 +114,14 @@ def validate_tfy_host(tfy_host: str) -> None:
 def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
     tfy_host = tfy_host or ENV_VARS.TFY_HOST
     if not tfy_host:
-        raise ValueError(
-            f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
-        )
+        if ENV_VARS.TFY_API_KEY:
+            raise ValueError(
+                f"TFY_HOST` env must be set since `{TFY_API_KEY_ENV_KEY}` env is set. Either set `{TFY_HOST_ENV_KEY}` or unset `{TFY_API_KEY_ENV_KEY}` and login"
+            )
+        else:
+            raise ValueError(
+                f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
+            )
     tfy_host = tfy_host.strip("/")
     validate_tfy_host(tfy_host)
     return tfy_host

truefoundry/deploy/__init__.py

@@ -120,6 +120,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     SparkJobPythonEntrypoint,
     SparkJobPythonNotebookEntrypoint,
     SparkJobScalaEntrypoint,
+    SparkJobScalaNotebookEntrypoint,
     SQSInputConfig,
     SQSOutputConfig,
     SQSQueueMetricConfig,

truefoundry/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-06-18T21:24:37+00:00
+#   timestamp: 2025-07-28T14:11:36+00:00
 
 from __future__ import annotations
 
@@ -103,7 +103,7 @@ class AsyncProcessorSidecar(BaseModel):
 
 
 class Autoshutdown(BaseModel):
-    wait_time: conint(ge=0) = Field(
+    wait_time: conint(ge=0, le=604800) = Field(
         900,
         description="The period to wait after the last received request before scaling the replicas to 0. This value should be high enough to allow for the replicas of the service to come up to avoid premature scaling down.",
     )
@@ -957,6 +957,14 @@ class SparkJobScalaEntrypoint(BaseModel):
     )
 
 
+class SparkJobScalaNotebookEntrypoint(BaseModel):
+    type: Literal["scala-notebook"] = Field(..., description="")
+    main_application_file: str = Field(
+        ...,
+        description="The main application file to be executed by the spark job. Relative path in case of git repository.",
+    )
+
+
 class StaticVolumeConfig(BaseModel):
     type: Literal["static"] = Field(..., description="Volume Type for the volume.")
     persistent_volume_name: str = Field(
@@ -989,6 +997,34 @@ class TaskDockerFileBuild(BaseModel):
     build_args: Optional[Dict[str, str]] = Field(None, description="")
 
 
+class TaskPySparkBuild(BaseModel):
+    """
+    Describes the configuration for the PySpark build for a task
+    """
+
+    type: Literal["task-pyspark-build"] = Field(..., description="")
+    spark_version: str = Field(
+        "3.5.2",
+        description="Spark version should match the spark version installed in the image.",
+    )
+    docker_registry: Optional[str] = Field(
+        None,
+        description="FQN of the container registry. If you can't find your registry here,\nadd it through the [Integrations](/integrations?tab=docker-registry) page",
+    )
+    requirements_path: Optional[str] = Field(
+        None,
+        description="Path to `requirements.txt` relative to\n`Path to build context`",
+    )
+    pip_packages: Optional[List[str]] = Field(
+        None,
+        description='Define pip package requirements.\nIn Python/YAML E.g. ["fastapi>=0.90,<1.0", "uvicorn"]',
+    )
+    apt_packages: Optional[List[str]] = Field(
+        None,
+        description='Debian packages to install via `apt get`.\nIn Python/YAML E.g. ["git", "ffmpeg", "htop"]',
+    )
+
+
 class TaskPythonBuild(BaseModel):
     """
     Describes the configuration for the python build for a task
@@ -1200,7 +1236,7 @@ class ContainerTaskConfig(BaseModel):
         description="Configure environment variables to be injected in the task either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
     )
     resources: Optional[Resources] = None
-    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
+    mounts: Optional[List[VolumeMount]] = Field(
         None, description="Configure data to be mounted to Workflow pod(s) as a volume."
     )
     service_account: Optional[str] = Field(None, description="")
@@ -1461,6 +1497,7 @@ class SparkJob(BaseModel):
         SparkJobScalaEntrypoint,
         SparkJobJavaEntrypoint,
         SparkJobPythonNotebookEntrypoint,
+        SparkJobScalaNotebookEntrypoint,
     ] = Field(..., description="")
     driver_config: SparkDriverConfig
     executor_config: SparkExecutorConfig
@@ -1567,15 +1604,6 @@ class FlyteLaunchPlan(BaseModel):
     closure: Any
 
 
-class FlyteTaskCustom(BaseModel):
-    truefoundry: Union[PythonTaskConfig, ContainerTaskConfig]
-
-
-class FlyteTaskTemplate(BaseModel):
-    id: FlyteTaskID
-    custom: FlyteTaskCustom
-
-
 class JobAlert(BaseModel):
     """
     Describes the configuration for the job alerts
@@ -1594,6 +1622,25 @@ class JobAlert(BaseModel):
     on_failure: bool = Field(True, description="Send an alert when the job fails")
 
 
+class PySparkTaskConfig(BaseModel):
+    type: Literal["pyspark-task-config"] = Field(..., description="")
+    image: TaskPySparkBuild
+    driver_config: SparkDriverConfig
+    executor_config: SparkExecutorConfig
+    spark_conf: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Extra configuration properties to be passed to the spark job. [Docs](https://spark.apache.org/docs/latest/configuration.html)",
+    )
+    env: Optional[Dict[str, str]] = Field(
+        None,
+        description="Configure environment variables to be injected in the task either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
+    )
+    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
+        None, description="Configure data to be mounted to Workflow pod(s) as a volume."
+    )
+    service_account: Optional[str] = Field(None, description="")
+
+
 class Service(BaseService):
     """
     Describes the configuration for the service
@@ -1630,9 +1677,13 @@ class AsyncService(BaseService):
     sidecar: Optional[AsyncProcessorSidecar] = None
 
 
-class FlyteTask(BaseModel):
-    template: FlyteTaskTemplate
-    description: Optional[Any] = None
+class FlyteTaskCustom(BaseModel):
+    truefoundry: Union[PythonTaskConfig, ContainerTaskConfig, PySparkTaskConfig]
+
+
+class FlyteTaskTemplate(BaseModel):
+    id: FlyteTaskID
+    custom: FlyteTaskCustom
 
 
 class Job(BaseModel):
@@ -1688,27 +1739,6 @@ class Job(BaseModel):
     )
 
 
-class Workflow(BaseModel):
-    """
-    Describes the configuration for the worflow
-    """
-
-    type: Literal["workflow"] = Field(..., description="")
-    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
-        ..., description="Name of the workflow"
-    )
-    source: Union[LocalSource, RemoteSource] = Field(
-        ..., description="Source Code for the workflow, either local or remote"
-    )
-    workflow_file_path: str = Field(
-        ..., description="Path to the workflow file relative to the project root path"
-    )
-    flyte_entities: Optional[List[Union[FlyteTask, FlyteWorkflow, FlyteLaunchPlan]]] = (
-        Field(None, description="")
-    )
-    alerts: Optional[List[WorkflowAlert]] = Field(None, description="")
-
-
 class ApplicationSet(BaseModel):
     """
     Describes the configuration for the application set
@@ -1735,6 +1765,32 @@ class ApplicationSet(BaseModel):
     )
 
 
+class FlyteTask(BaseModel):
+    template: FlyteTaskTemplate
+    description: Optional[Any] = None
+
+
+class Workflow(BaseModel):
+    """
+    Describes the configuration for the worflow
+    """
+
+    type: Literal["workflow"] = Field(..., description="")
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
+        ..., description="Name of the workflow"
+    )
+    source: Union[LocalSource, RemoteSource] = Field(
+        ..., description="Source Code for the workflow, either local or remote"
+    )
+    workflow_file_path: str = Field(
+        ..., description="Path to the workflow file relative to the project root path"
+    )
+    flyte_entities: Optional[List[Union[FlyteTask, FlyteWorkflow, FlyteLaunchPlan]]] = (
+        Field(None, description="")
+    )
+    alerts: Optional[List[WorkflowAlert]] = Field(None, description="")
+
+
 class Application(BaseModel):
     __root__: Union[
         Service,

truefoundry/deploy/builder/__init__.py

@@ -5,6 +5,7 @@ from truefoundry.deploy._autogen.models import (
     PythonBuild,
     SparkBuild,
     TaskDockerFileBuild,
+    TaskPySparkBuild,
     TaskPythonBuild,
 )
 from truefoundry.deploy.builder.builders import get_builder
@@ -22,6 +23,7 @@ class _BuildConfig(BaseModel):
         TaskPythonBuild,
         TaskDockerFileBuild,
         SparkBuild,
+        TaskPySparkBuild,
     ] = Field(discriminator="type")
 
 

truefoundry/deploy/builder/builders/__init__.py

@@ -5,6 +5,7 @@ from truefoundry.deploy.builder.builders import (
     tfy_notebook_buildpack,
     tfy_python_buildpack,
     tfy_spark_buildpack,
+    tfy_task_pyspark_buildpack,
 )
 
 BUILD_REGISTRY: Dict[str, Callable] = {
@@ -12,6 +13,7 @@ BUILD_REGISTRY: Dict[str, Callable] = {
     "tfy-python-buildpack": tfy_python_buildpack.build,
     "tfy-notebook-buildpack": tfy_notebook_buildpack.build,
     "tfy-spark-buildpack": tfy_spark_buildpack.build,
+    "task-pyspark-build": tfy_task_pyspark_buildpack.build,
 }
 
 __all__ = ["get_builder"]

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -1,5 +1,4 @@
-import shlex
-from typing import Dict, List, Optional
+from typing import Dict
 
 from mako.template import Template
 
@@ -7,9 +6,12 @@ from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy._autogen.models import PythonBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
-    PIP_CONF_SECRET_MOUNT_AS_ENV,
     UV_CONF_BUILDKIT_SECRET_MOUNT,
-    UV_CONF_SECRET_MOUNT_AS_ENV,
+)
+from truefoundry.deploy.builder.utils import (
+    generate_apt_install_command,
+    generate_pip_install_command,
+    generate_uv_pip_install_command,
 )
 from truefoundry.deploy.v2.lib.patched_models import (
     CUDAVersion,
@@ -82,85 +84,6 @@ CUDA_VERSION_TO_IMAGE_TAG: Dict[str, str] = {
 }
 
 
-def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[str]:
-    packages_list = None
-    if apt_packages:
-        packages_list = " ".join(p.strip() for p in apt_packages if p.strip())
-    if not packages_list:
-        return None
-    apt_update_command = "apt update"
-    apt_install_command = f"DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends {packages_list}"
-    clear_apt_lists_command = "rm -rf /var/lib/apt/lists/*"
-    return " && ".join(
-        [apt_update_command, apt_install_command, clear_apt_lists_command]
-    )
-
-
-def generate_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_pip_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
-    envs = []
-    if mount_pip_conf_secret:
-        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["python", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
-    args = []
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    final_pip_install_command = shlex.join(envs + command + args)
-    final_docker_run_command = " && ".join(
-        [upgrade_pip_command, final_pip_install_command]
-    )
-    return final_docker_run_command
-
-
-def generate_uv_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_uv_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
-    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
-    envs = [
-        "UV_LINK_MODE=copy",
-        "UV_PYTHON_DOWNLOADS=never",
-        "UV_INDEX_STRATEGY=unsafe-best-match",
-    ]
-    if mount_uv_conf_secret:
-        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["uv", "pip", "install", "--no-cache-dir"]
-
-    args = []
-
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    uv_pip_install_command = shlex.join(envs + command + args)
-    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
-    final_docker_run_command = " ".join([uv_mount, shell_commands])
-
-    return final_docker_run_command
-
-
 def generate_dockerfile_content(
     build_configuration: PythonBuild,
     package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,

truefoundry/deploy/builder/builders/tfy_spark_buildpack/dockerfile_template.py

@@ -1,15 +1,14 @@
-import shlex
-from typing import List, Optional
-
 from mako.template import Template
 
 from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy._autogen.models import SparkBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
-    PIP_CONF_SECRET_MOUNT_AS_ENV,
     UV_CONF_BUILDKIT_SECRET_MOUNT,
-    UV_CONF_SECRET_MOUNT_AS_ENV,
+)
+from truefoundry.deploy.builder.utils import (
+    generate_pip_install_command,
+    generate_uv_pip_install_command,
 )
 from truefoundry.deploy.v2.lib.patched_models import (
     _resolve_requirements_path,
@@ -26,8 +25,7 @@ RUN ${package_manager_config_secret_mount} ${python_packages_install_command}
 % endif
 ENV PYTHONDONTWRITEBYTECODE=1
 ENV IPYTHONDIR=/tmp/.ipython
-RUN groupadd --system --gid 1001 spark && useradd --system --uid 1001 --gid spark --no-create-home spark
-USER spark
+USER 1001
 COPY . /app
 """
 
@@ -35,6 +33,20 @@ _POST_USER_TEMPLATE = """
 COPY tfy_execute_notebook.py /app/tfy_execute_notebook.py
 """
 
+_ALMOND_INSTALL_TEMPLATE = """
+ENV COURSIER_CACHE=/opt/coursier-cache
+RUN install_packages curl
+RUN curl -Lo coursier https://git.io/coursier-cli && \
+    chmod +x coursier && \
+    ./coursier launch almond:0.14.1 -- --install --global && \
+    chown -R 1001:0 /usr/local/share/jupyter && \
+    chown -R 1001:0 /opt/coursier-cache && \
+    rm -f coursier
+"""
+
+# Docker image size with almond - 1.26GB
+# Docker image size without almond - 1.1GB
+# Not much harm in packaging almond by default
 DOCKERFILE_TEMPLATE = Template(
     """
 FROM ${spark_image_repo}:${spark_version}
@@ -43,6 +55,7 @@ RUN apt update && \
     DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git && \
     rm -rf /var/lib/apt/lists/*
 """
+    + _ALMOND_INSTALL_TEMPLATE
     + _POST_PYTHON_INSTALL_TEMPLATE
     + _POST_USER_TEMPLATE
 )
@@ -55,71 +68,6 @@ ADDITIONAL_PIP_PACKAGES = [
 ]
 
 
-def generate_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_pip_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python3 -m pip install -U pip setuptools wheel"
-    envs = []
-    if mount_pip_conf_secret:
-        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["python3", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
-    args = []
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    final_pip_install_command = shlex.join(envs + command + args)
-    final_docker_run_command = " && ".join(
-        [upgrade_pip_command, final_pip_install_command]
-    )
-    return final_docker_run_command
-
-
-def generate_uv_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_uv_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python3 -m pip install -U pip setuptools wheel"
-    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
-    envs = [
-        "UV_LINK_MODE=copy",
-        "UV_PYTHON_DOWNLOADS=never",
-        "UV_INDEX_STRATEGY=unsafe-best-match",
-    ]
-    if mount_uv_conf_secret:
-        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["uv", "pip", "install", "--no-cache-dir"]
-
-    args = []
-
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    uv_pip_install_command = shlex.join(envs + command + args)
-    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
-    final_docker_run_command = " ".join([uv_mount, shell_commands])
-
-    return final_docker_run_command
-
-
 def generate_dockerfile_content(
     build_configuration: SparkBuild,
     package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,

truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py

@@ -45,16 +45,9 @@ def upload_file_to_s3(file_path, bucket_name, s3_key):
     # Use s3proxy for pushing data to s3
     # The JWT token is already available in the pod
     aws_access_key_id = os.environ.get("SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN")
-    aws_secret_access_key = "__token__"
+    aws_secret_access_key = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_SECRET_KEY")
     s3_endpoint_url = os.environ.get("S3_PROXY_URL")
 
-    if not aws_access_key_id:
-        raise ValueError(
-            "SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN environment variable is not set"
-        )
-    if not s3_endpoint_url:
-        raise ValueError("S3_PROXY_URL environment variable is not set")
-
     # Needed for the issue https://github.com/gaul/s3proxy/issues/765
     s3_config = Config(
         request_checksum_calculation="when_required",
@@ -93,12 +86,24 @@ def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=
         parameters = {}
 
     print(f"Starting execution of notebook: {notebook_path}")
+    notebook_type = os.environ.get("TFY_NOTEBOOK_TYPE", "").lower()
+    kernel_mapping = {"python": "python3", "scala": "scala"}
+
+    if notebook_type not in kernel_mapping:
+        supported_types = ", ".join(kernel_mapping.keys())
+        raise ValueError(
+            f"Unsupported notebook type: '{notebook_type}'. "
+            f"Supported types: [{supported_types}]"
+        )
+
+    kernel_name = kernel_mapping[notebook_type]
+
     pm.execute_notebook(
         input_path=notebook_path,
         output_path=output_path,
         parameters=parameters,
         # TODO(gw): Replace with kernel name for venv
-        kernel_name="python3",
+        kernel_name=kernel_name,
         # Log cell by cell execution output
         # TODO(gw): Output logs to a file instead, so that they aren't merged with the container's logs
         log_output=True,
@@ -106,7 +111,23 @@ def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=
         stderr_file=sys.stderr,
     )
     print(f"Successfully executed notebook: {notebook_path}")
-    return output_path
+
+
+def validate_env_vars():
+    keys = [
+        "TFY_NOTEBOOK_OUTPUT_S3_KEY",
+        "TFY_NOTEBOOK_OUTPUT_S3_BUCKET",
+        "SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN",
+        "TFY_NOTEBOOK_OUTPUT_S3_SECRET_KEY",
+        "TFY_NOTEBOOK_TYPE",
+    ]
+    unset_keys = [key for key in keys if not os.environ.get(key)]
+    if unset_keys:
+        raise ValueError(
+            f"Environment variables {unset_keys} are not set."
+            f"Contact you tenant-admin to configure storage bucket on the control plane "
+            f"to enable uploading spark notebook outputs."
+        )
 
 
 if __name__ == "__main__":
@@ -116,35 +137,31 @@ if __name__ == "__main__":
     parser.add_argument("notebook_path", help="Path to the notebook file to execute")
     args = parser.parse_args()
 
+    # Since failure to upload is considered a job failure, fail the job even before it run if uploads cannot happen
+    validate_env_vars()
+
     output_notebook_path = "/tmp/output.ipynb"
 
     # This would be the same as the default bucket used by servicefoundry-server
     s3_bucket = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_BUCKET")
-    # This would be something like sparkjob-events/<tenant-id>
-    s3_key_prefix = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_KEY_PREFIX")
+    # This would be something like sparkjob-events/<tenant-id>/output-notebooks/<application-id>/<jobrun-name>/output.html
+    s3_key = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_KEY")
 
     try:
-        executed_notebook_path = execute_notebook(
-            args.notebook_path, output_path=output_notebook_path
-        )
+        execute_notebook(args.notebook_path, output_path=output_notebook_path)
 
         # The following may also be modeled as an entrypoint
         # https://papermill.readthedocs.io/en/latest/extending-entry-points.html
         # Will take that up with next iteration where we save the executed notebook periodically
-        if s3_bucket and s3_key_prefix:
-            print("Converting notebook to HTML and uploading to S3...")
-            html_output_path = "/tmp/output.html"
-            convert_notebook_to_html(
-                notebook_path=executed_notebook_path, output_html_path=html_output_path
-            )
-
-            # Construct S3 key: use the original notebook name for the HTML file
-            notebook_name = os.path.basename(args.notebook_path)
-            s3_html_key = f"{s3_key_prefix}/output.html"
-            upload_file_to_s3(
-                file_path=html_output_path, bucket_name=s3_bucket, s3_key=s3_html_key
-            )
-            print(f"Successfully uploaded HTML to s3://{s3_bucket}/{s3_html_key}")
+        print("Converting notebook to HTML and uploading to S3...")
+        html_output_path = "/tmp/output.html"
+        convert_notebook_to_html(
+            notebook_path=output_notebook_path, output_html_path=html_output_path
+        )
+        upload_file_to_s3(
+            file_path=html_output_path, bucket_name=s3_bucket, s3_key=s3_key
+        )
+        print(f"Successfully uploaded HTML to s3://{s3_bucket}/{s3_key}")
 
     except Exception as e:
         print(f"Error executing notebook {args.notebook_path}: {e}")