truefoundry 0.2.10__py3-none-any.whl → 0.3.0__py3-none-any.whl

truefoundry/workflow/example/deploy.sh

@@ -0,0 +1 @@
+ tfy deploy --workspace_fqn tfy-ctl-euwe1-devtest:nikhil-ws

truefoundry/workflow/example/hello_world_package/workflow.py

@@ -0,0 +1,20 @@
+from truefoundry.workflow import PythonTaskConfig, TaskPythonBuild, task, workflow
+
+
+@task(
+    task_config=PythonTaskConfig(
+        image=TaskPythonBuild(
+            python_version="3.9",
+            pip_packages=["truefoundry[workflow]==0.3.0rc7"],
+        ),
+        service_account="tfy-flyte-dataplane-devtest-s3",
+    )
+)
+def say_hello() -> str:
+    return "Hello, World!"
+
+
+@workflow
+def hello_world_wf() -> str:
+    res = say_hello()
+    return res

truefoundry/workflow/example/package/test_workflow.py

@@ -0,0 +1,152 @@
+import os
+from functools import partial
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from truefoundry.deploy import Image, NvidiaGPU, Resources
+from truefoundry.workflow import (
+    ContainerTask,
+    ContainerTaskConfig,
+    ExecutionConfig,
+    FlyteDirectory,
+    PythonTaskConfig,
+    TaskPythonBuild,
+    conditional,
+    map_task,
+    task,
+    workflow,
+)
+
+cpu_task_config = PythonTaskConfig(
+    image=TaskPythonBuild(
+        python_version="3.9",
+        pip_packages=["truefoundry[workflow]==0.3.0rc7"],
+    ),
+    resources=Resources(cpu_request=0.45),
+    service_account="tfy-flyte-dataplane-devtest-s3",
+)
+
+
+# figure out naming
+@task(task_config=cpu_task_config)
+def should_train_tokenizer(tokenizer: str) -> bool:
+    print("Should train tokenizer")
+    return not bool(tokenizer)
+
+
+@task(
+    task_config=cpu_task_config,
+    # currently we will not support the caching, but keeping here for now
+    # cache=True,
+    # cache_version="1.0",
+)
+def train_tokenizer() -> str:
+    print("Training tokenizer")
+    return "trained_tokenizer"
+
+
+@task(
+    task_config=PythonTaskConfig(
+        image=TaskPythonBuild(
+            python_version="3.9",
+            pip_packages=["truefoundry[workflow]==0.3.0rc7", "pynvml==11.5.0"],
+            cuda_version="11.5-cudnn8",
+        ),
+        env={
+            "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
+            "NVIDIA_VISIBLE_DEVICES": "all",
+        },
+        resources=Resources(cpu_request=0.45, devices=[NvidiaGPU(name="T4", count=1)]),
+        service_account="tfy-flyte-dataplane-devtest-s3",
+    ),
+)
+def train_model(tokenizer: str) -> Tuple[FlyteDirectory, str]:
+    print("Training model")
+    import flytekit
+    from pynvml import nvmlDeviceGetCount, nvmlInit
+
+    nvmlInit()
+    assert nvmlDeviceGetCount() > 0
+
+    working_dir = flytekit.current_context().working_directory
+    local_dir = Path(os.path.join(working_dir, "csv_files"))
+    local_dir.mkdir(exist_ok=True)
+
+    with open(os.path.join(local_dir, "model"), "w", encoding="utf-8") as f:
+        f.write(tokenizer)
+
+    return FlyteDirectory(path=str(local_dir)), "hello"
+
+
+@task(task_config=cpu_task_config)
+def get_validation_data() -> List[str]:
+    print("Getting validation data")
+    return ["foo", "bar", "baz"]
+
+
+@task(task_config=cpu_task_config)
+def validate_model(model: FlyteDirectory, tokenizer: str, validation_data: str) -> bool:
+    print(validation_data)
+    model_path = os.path.join(model, "model")
+    with open(model_path, "r", encoding="utf-8") as f:
+        return f.read() == tokenizer
+
+
+@task(task_config=cpu_task_config)
+def all_good(validations: List[bool]) -> bool:
+    print("Validations", validations)
+    return all(validations)
+
+
+echo = ContainerTask(
+    name="echo",
+    task_config=ContainerTaskConfig(
+        image=Image(
+            image_uri="bash:4.1",
+            command=["echo", "hello"],
+        ),
+        service_account="tfy-flyte-dataplane-devtest-s3",
+    ),
+)
+
+
+@task(task_config=cpu_task_config)
+def random(tokenizer: str) -> Tuple[FlyteDirectory, str]:
+    print(tokenizer)
+    return FlyteDirectory(path=""), "random"
+
+
+@workflow(
+    execution_configs=[
+        ExecutionConfig(
+            schedule="*/10 * * * *",
+        )
+    ]
+)
+def train(
+    tokenizer: str = "",
+    test_v2: str = "",
+    optional_var: Optional[str] = "",
+    default_v: Optional[str] = "hello1",
+    default_v2: str = "hello2",
+) -> bool:
+    stt = should_train_tokenizer(tokenizer=tokenizer)
+    model, t = (
+        conditional("train_tokenizer")
+        .if_(stt.is_true())
+        .then(train_model(tokenizer=tokenizer))
+        .else_()
+        .then(random(tokenizer=tokenizer))
+    )
+    validation_task = partial(validate_model, model=model, tokenizer=t)
+    validation_data = get_validation_data()
+    validations = map_task(
+        validation_task,
+        concurrency=2,
+    )(validation_data=validation_data)
+    echo()
+    return all_good(validations=validations)
+
+
+if __name__ == "__main__":
+    train()

truefoundry/workflow/example/truefoundry.yaml

@@ -0,0 +1,9 @@
+type: workflow
+name: my-test-workflow9
+source:
+  type: local
+  local_build: false
+  project_root_path: "./"
+workflow_file_path: package/test_workflow.py
+# workflow_file_path: hello_world_package/workflow.py
+flyte_entities: []

truefoundry/workflow/example/workflow.yaml

@@ -0,0 +1,116 @@
+type: workflow
+name: my-test-workflow5
+source:
+  type: remote
+  remote_uri: s3://tfy-mlfoundry-tfy-ctl-euwe1-devtest20221012173221749100000002/tfy-ctl-euwe1-devtest:nikhil-ws/my-test-workflow5/clyhb79w4013l01oyfq1iawd5/package.tar.gz
+workflow_file_path: hello_world_package/workflow.py
+docker_registry: docker.io/nikp-test
+flyte_entities:
+- template:
+    id:
+      resourceType: TASK
+      name: hello_world_package.workflow.say_hello
+    type: python-task
+    metadata:
+      runtime:
+        type: FLYTE_SDK
+        version: 1.12.2
+        flavor: python
+      retries: {}
+    interface:
+      inputs: {}
+      outputs:
+        variables:
+          o0:
+            type:
+              simple: STRING
+            description: o0
+    custom:
+      truefoundry:
+        type: python-task-config
+        image:
+          type: task-python-build
+          python_version: '3.9'
+          pip_packages:
+          - flytekit==1.10.3
+        resources:
+          cpu_request: 0.2
+          cpu_limit: 0.5
+          memory_request: 200
+          memory_limit: 500
+          ephemeral_storage_request: 1000
+          ephemeral_storage_limit: 2000
+    container:
+      image: ':'
+      args:
+      - pyflyte-execute
+      - --inputs
+      - '{{.input}}'
+      - --output-prefix
+      - '{{.outputPrefix}}'
+      - --raw-output-data-prefix
+      - '{{.rawOutputDataPrefix}}'
+      - --checkpoint-path
+      - '{{.checkpointOutputPrefix}}'
+      - --prev-checkpoint
+      - '{{.prevCheckpointPrefix}}'
+      - --resolver
+      - flytekit.core.python_auto_container.default_task_resolver
+      - --
+      - task-module
+      - hello_world_package.workflow
+      - task-name
+      - say_hello
+      resources: {}
+  description:
+    longDescription:
+      format: DESCRIPTION_FORMAT_RST
+- template:
+    id:
+      resourceType: WORKFLOW
+      name: hello_world_package.workflow.hello_world_wf
+    metadata: {}
+    interface:
+      inputs: {}
+      outputs:
+        variables:
+          o0:
+            type:
+              simple: STRING
+            description: o0
+    nodes:
+    - id: n0
+      metadata:
+        name: say_hello
+        retries: {}
+      taskNode:
+        referenceId:
+          resourceType: TASK
+          name: hello_world_package.workflow.say_hello
+        overrides: {}
+    outputs:
+    - var: o0
+      binding:
+        promise:
+          nodeId: n0
+          var: o0
+    metadataDefaults: {}
+  description:
+    longDescription:
+      format: DESCRIPTION_FORMAT_RST
+- id:
+    resourceType: LAUNCH_PLAN
+    name: hello_world_package.workflow.hello_world_wf
+  spec:
+    workflowId:
+      resourceType: WORKFLOW
+      name: hello_world_package.workflow.hello_world_wf
+    entityMetadata: {}
+    defaultInputs: {}
+    fixedInputs: {}
+    labels: {}
+    annotations: {}
+    rawOutputDataConfig: {}
+  closure:
+    expectedInputs: {}
+    expectedOutputs: {}

truefoundry/workflow/map_task.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+from flytekit import PythonFunctionTask
+from flytekit.core.array_node_map_task import ArrayNodeMapTask
+
+
+class TrueFoundryArrayNodeMapTask(ArrayNodeMapTask):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_custom(self, settings) -> dict:
+        arraynode_custom_params = super().get_custom(settings)
+        return {
+            "truefoundry": self._partial.func.task_config.dict(),
+            **arraynode_custom_params,
+        }
+
+
+def map_task(
+    task_function: PythonFunctionTask,
+    concurrency: Optional[int] = None,
+    # TODO why no min_successes?
+    min_success_ratio: float = 1.0,
+    **kwargs,
+):
+    """Map task that uses the ``ArrayNode`` construct..
+
+    .. important::
+
+       This is an experimental drop-in replacement for :py:func:`~flytekit.map_task`.
+
+    :param task_function: This argument is implicitly passed and represents the repeatable function
+    :param concurrency: If specified, this limits the number of mapped tasks than can run in parallel to the given batch
+        size. If the size of the input exceeds the concurrency value, then multiple batches will be run serially until
+        all inputs are processed. If set to 0, this means unbounded concurrency. If left unspecified, this means the
+        array node will inherit parallelism from the workflow
+    :param min_success_ratio: If specified, this determines the minimum fraction of total jobs which can complete
+        successfully before terminating this task and marking it successful.
+    """
+    return TrueFoundryArrayNodeMapTask(
+        task_function,
+        concurrency=concurrency,
+        min_success_ratio=min_success_ratio,
+        **kwargs,
+    )

truefoundry/workflow/python_task.py

@@ -0,0 +1,32 @@
+from typing import Callable, Tuple, Union
+
+from flytekit import FlyteContext, PythonFunctionTask
+from flytekit.extend import Promise, TaskPlugins
+
+from truefoundry.logger import logger
+from truefoundry.workflow import PythonTaskConfig
+
+
+class TrueFoundryFunctionTask(PythonFunctionTask[PythonTaskConfig]):
+    # Note(nikp1172) We are not creating our task type for now, we just re-use their task type for now
+    def __init__(
+        self,
+        task_config: PythonTaskConfig,
+        task_function: Callable,
+        **kwargs,
+    ):
+        super().__init__(task_config=task_config, task_function=task_function, **kwargs)
+
+    def local_execute(
+        self, ctx: FlyteContext, **kwargs
+    ) -> Union[Tuple[Promise], Promise, None]:
+        logger.warning(
+            "Running pod task locally. Local environment may not match pod environment which may cause issues."
+        )
+        return super().local_execute(ctx=ctx, **kwargs)
+
+    def get_custom(self, settings) -> dict:
+        return {"truefoundry": self._task_config.dict()}
+
+
+TaskPlugins.register_pythontask_plugin(PythonTaskConfig, TrueFoundryFunctionTask)

truefoundry/workflow/task.py

@@ -0,0 +1,50 @@
+from typing import Any, Callable, Optional, Union
+
+from flytekit import PythonFunctionTask
+from flytekit.core.task import FuncOut, T
+from flytekit.core.task import task as flytekit_task
+
+from truefoundry.workflow import PythonTaskConfig
+
+
+def task(
+    _task_function: Optional[Callable[..., Any]] = None,
+    task_config: Optional[PythonTaskConfig] = None,
+    retries: Optional[int] = 0,
+) -> Union[
+    Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]],
+    PythonFunctionTask[T],
+    Callable[..., FuncOut],
+]:
+    """
+    This is the decorator used to run flyte tasks in Truefoundry.
+
+    Tasks are the building blocks of workflow. They represent users code. Tasks have the following properties
+
+    * Versioned (usually tied to the git revision SHA1)
+    * Strong interfaces (specified inputs and outputs)
+    * Declarative
+    * Independently executable
+    * Unit testable
+
+    For a simple python task,
+
+    .. code-block:: python
+        from truefoundry.workflow import task
+        @task
+        def my_task(x: int, y: typing.Dict[str, str]) -> str:
+            ...
+
+    For specific task types
+
+    .. code-block:: python
+        from truefoundry.workflow import task, PythonTaskConfig
+        @task(task_config=PythonTaskConfig(), retries=2)
+        def my_task(x: int, y: typing.Dict[str, str]) -> str:
+            ...
+
+    :param _task_function: The function to be wrapped as a task. This is optional and can be used as a decorator.
+    :param task_config: The configuration for the task of class PythonTaskConfig. This is optional and defaults to an empty configuration.
+    :param retries: The number of retries for the task. This is optional and defaults to 0.
+    """
+    return flytekit_task(_task_function, task_config=task_config, retries=retries)

truefoundry/workflow/workflow.py

@@ -0,0 +1,114 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from flytekit.core.base_task import Task
+from flytekit.core.launch_plan import LaunchPlan
+from flytekit.core.schedule import CronSchedule
+from flytekit.core.workflow import (
+    FuncOut,
+    PythonFunctionWorkflow,
+    WorkflowFailurePolicy,
+)
+from flytekit.core.workflow import workflow as flytekit_workflow
+
+from truefoundry.pydantic_v1 import BaseModel
+
+TRUEFOUNDRY_LAUNCH_PLAN_NAME = "default"
+
+
+class ExecutionConfig(BaseModel):
+    # I am not using job's Schedule here because:
+    # 1. flyte accepts schedule only as a string
+    # 2. flyte doesn't support setting concurrency policy
+    schedule: str
+
+
+class ExecutionConfigStore(BaseModel):
+    execution_config_map: Dict[str, ExecutionConfig] = {}
+
+    def reset(self):
+        self.execution_config_map = {}
+
+    def get(self, key: str) -> Optional[ExecutionConfig]:
+        return self.execution_config_map.get(key)
+
+    def set(self, key: str, value: ExecutionConfig):
+        self.execution_config_map[key] = value
+
+
+execution_config_store = ExecutionConfigStore()
+
+
+def workflow(
+    _workflow_function: Optional[Callable[..., Any]] = None,
+    failure_policy: Optional[WorkflowFailurePolicy] = None,
+    on_failure: Optional[Task] = None,
+    execution_configs: Optional[List[ExecutionConfig]] = None,
+) -> Union[
+    Callable[[Callable[..., FuncOut]], PythonFunctionWorkflow],
+    PythonFunctionWorkflow,
+    Callable[..., FuncOut],
+]:
+    """
+    This decorator declares a function to be a Flyte workflow. Workflows are declarative entities that construct a DAG
+    of tasks using the data flow between tasks.
+
+    Unlike a task, the function body of a workflow is evaluated at serialization-time (aka compile-time). This is
+    because while we can determine the entire structure of a task by looking at the function's signature, workflows need
+    to run through the function itself because the body of the function is what expresses the workflow structure. It's
+    also important to note that, local execution notwithstanding, it is not evaluated again when the workflow runs on
+    Flyte.
+    That is, workflows should not call non-Flyte entities since they are only run once (again, this is with respect to
+    the platform, local runs notwithstanding).
+
+    Example:
+
+    .. literalinclude:: ../../../tests/flytekit/unit/core/test_workflows.py
+       :pyobject: my_wf_example
+
+    Again, users should keep in mind that even though the body of the function looks like regular Python, it is
+    actually not. When flytekit scans the workflow function, the objects being passed around between the tasks are not
+    your typical Python values. So even though you may have a task ``t1() -> int``, when ``a = t1()`` is called, ``a``
+    will not be an integer so if you try to ``range(a)`` you'll get an error.
+
+    Please see the :ref:`user guide <cookbook:workflow>` for more usage examples.
+
+    :param _workflow_function: This argument is implicitly passed and represents the decorated function.
+    :param failure_policy: Use the options in flytekit.WorkflowFailurePolicy
+    :param on_failure: Invoke this workflow or task on failure. The Workflow / task has to match the signature of
+         the current workflow, with an additional parameter called `error` Error
+    """
+    if _workflow_function is None:
+
+        def wrapper(func: Callable[..., Any]) -> Any:
+            return workflow(
+                _workflow_function=func,
+                failure_policy=failure_policy,
+                on_failure=on_failure,
+                execution_configs=execution_configs,
+            )
+
+        return wrapper
+
+    workflow_object = flytekit_workflow(
+        _workflow_function,
+        failure_policy=failure_policy,
+        on_failure=on_failure,
+    )
+    execution_configs = execution_configs or []
+    if len(execution_configs) > 1:
+        raise ValueError("Only one cron execution config is allowed per workflow")
+
+    if execution_configs:
+        execution_config = execution_configs[0]
+        # flyte maintains this in-memory and uses it while serialization
+        LaunchPlan.get_or_create(
+            workflow=workflow_object,
+            name=TRUEFOUNDRY_LAUNCH_PLAN_NAME,
+            schedule=CronSchedule(schedule=execution_config.schedule),
+        )
+        function_module = _workflow_function.__module__
+        function_name = _workflow_function.__name__
+        function_path = f"{function_module}.{function_name}"
+        execution_config_store.set(function_path, execution_config)
+
+    return workflow_object

{truefoundry-0.2.10.dist-info → truefoundry-0.3.0.dist-info}/METADATA

@@ -1,27 +1,47 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.2.10
+Version: 0.3.0
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
-Requires-Python: >=3.8,<3.13
+Requires-Python: >=3.8.1,<3.13
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: ml
-Requires-Dist: docker (>=7.0.0,<8.0.0)
+Provides-Extra: workflow
+Requires-Dist: GitPython (>=3.1.43,<3.2.0)
+Requires-Dist: Mako (>=1.1.6,<2.0.0)
+Requires-Dist: PyJWT (>=2.0.0,<3.0.0)
+Requires-Dist: PyYAML (>=6.0.0,<7.0.0)
+Requires-Dist: click (>=7.0.0,<9.0.0)
+Requires-Dist: docker (>=6.1.2,<8.0.0)
+Requires-Dist: fastapi (>=0.56.0,<0.200.0)
+Requires-Dist: filelock (>=3.8.0,<4.0.0)
+Requires-Dist: flytekit (==1.12.2) ; extra == "workflow"
 Requires-Dist: gitignorefile (>=1.1.2,<1.2.0)
-Requires-Dist: gitpython (>=3.1.43,<3.2.0)
-Requires-Dist: mlfoundry (==0.11.2) ; extra == "ml"
+Requires-Dist: importlib-metadata (>=6.0.1,<8.0.0)
+Requires-Dist: importlib-resources (>=5.2.0,<6.0.0)
+Requires-Dist: mlfoundry (==0.11.5) ; extra == "ml"
 Requires-Dist: openai (>=1.16.2,<2.0.0)
+Requires-Dist: packaging (>=20.0,<25.0)
 Requires-Dist: pydantic (>=1.10.0,<3)
+Requires-Dist: pygments (>=2.12.0,<3.0.0)
+Requires-Dist: python-dateutil (>=2.8.2,<3.0.0)
 Requires-Dist: python-dotenv (>=1.0.1,<1.1.0)
+Requires-Dist: python-socketio[client] (>=5.5.2,<6.0.0)
+Requires-Dist: questionary (>=1.10.0,<2.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: requirements-parser (>=0.10.2,<0.11.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
-Requires-Dist: servicefoundry (==0.10.16)
+Requires-Dist: rich-click (>=1.2.1,<2.0.0)
+Requires-Dist: tqdm (>=4.0.0,<5.0.0)
+Requires-Dist: typing-extensions (>=3.10.0)
+Requires-Dist: urllib3 (>=1.26.18,<3)
+Requires-Dist: uvicorn (>=0.13.0,<0.40.0)
+Requires-Dist: yq (>=3.1.0,<4.0.0)
 Description-Content-Type: text/markdown
 
 # TrueFoundry