flytekitplugins-perian-job 1.12.0__tar.gz

flytekitplugins-perian_job-1.12.0/PKG-INFO

@@ -0,0 +1,92 @@
+Metadata-Version: 2.1
+Name: flytekitplugins-perian_job
+Version: 1.12.0
+Summary: Flyte agent for Perian Job Platform (perian.io)
+Author: Omar Tarabai
+Author-email: otarabai@perian.io
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+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: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+
+# Flytekit Perian Job Platform Plugin
+
+Flyte Agent plugin for executing Flyte tasks on Perian Job Platform (perian.io).
+
+Perian Job Platform is still in closed beta. Contact support@perian.io if you are interested in trying it out.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-perian-job
+```
+
+## Getting Started
+
+This plugin allows executing `PythonFunctionTask` on Perian.
+
+An [ImageSpec](https://docs.flyte.org/en/latest/user_guide/customizing_dependencies/imagespec.html) need to be built with the perian agent plugin installed.
+
+### Parameters
+
+The following parameters can be used to set the requirements for the Perian task. If any of the requirements are skipped, it is replaced with the cheapest option. At least one requirement value should be set.
+* `cores`: Number of CPU cores
+* `memory`: Amount of memory in GB
+* `accelerators`: Number of accelerators
+* `accelerator_type`: Type of accelerator (e.g. 'A100'). For a full list of supported accelerators, use the perian CLI list-accelerators command.
+* `country_code`: Country code to run the job in (e.g. 'DE')
+
+### Credentials
+
+The following [secrets](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html) are required to be defined for the agent server:
+* Perian credentials:
+    * `perian_organization`
+    * `perian_token`
+* AWS credentials for accessing the Flyte storage bucket. You might need to create long-lived credentials to avoid expiry when running tasks. These credentials are never logged by Perian and are stored only until it is used, then immediately deleted:
+    * `aws_access_key_id`
+    * `aws_secret_access_key`
+* (Optional) Custom docker registry for pulling the Flyte image:
+    * `docker_registry_url`
+    * `docker_registry_username`
+    * `docker_registry_password`
+
+### Example
+
+`example.py` workflow example:
+```python
+from flytekit import ImageSpec, task, workflow
+from flytekitplugins.perian_job import PerianConfig
+
+image_spec = ImageSpec(
+    name="flyte-test",
+    registry="my-registry",
+    python_version="3.11",
+    apt_packages=["wget", "curl", "git"],
+    packages=[
+        "flytekitplugins-perian-job",
+    ],
+)
+
+@task(container_image=image_spec,
+      task_config=PerianConfig(
+        accelerators=1,
+        accelerator_type="A100",
+      ))
+def perian_hello(name: str) -> str:
+    return f"hello {name}!"
+
+@workflow
+def my_wf(name: str = "world") -> str:
+    return perian_hello(name=name)
+```

flytekitplugins-perian_job-1.12.0/README.md

@@ -0,0 +1,70 @@
+# Flytekit Perian Job Platform Plugin
+
+Flyte Agent plugin for executing Flyte tasks on Perian Job Platform (perian.io).
+
+Perian Job Platform is still in closed beta. Contact support@perian.io if you are interested in trying it out.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-perian-job
+```
+
+## Getting Started
+
+This plugin allows executing `PythonFunctionTask` on Perian.
+
+An [ImageSpec](https://docs.flyte.org/en/latest/user_guide/customizing_dependencies/imagespec.html) need to be built with the perian agent plugin installed.
+
+### Parameters
+
+The following parameters can be used to set the requirements for the Perian task. If any of the requirements are skipped, it is replaced with the cheapest option. At least one requirement value should be set.
+* `cores`: Number of CPU cores
+* `memory`: Amount of memory in GB
+* `accelerators`: Number of accelerators
+* `accelerator_type`: Type of accelerator (e.g. 'A100'). For a full list of supported accelerators, use the perian CLI list-accelerators command.
+* `country_code`: Country code to run the job in (e.g. 'DE')
+
+### Credentials
+
+The following [secrets](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html) are required to be defined for the agent server:
+* Perian credentials:
+    * `perian_organization`
+    * `perian_token`
+* AWS credentials for accessing the Flyte storage bucket. You might need to create long-lived credentials to avoid expiry when running tasks. These credentials are never logged by Perian and are stored only until it is used, then immediately deleted:
+    * `aws_access_key_id`
+    * `aws_secret_access_key`
+* (Optional) Custom docker registry for pulling the Flyte image:
+    * `docker_registry_url`
+    * `docker_registry_username`
+    * `docker_registry_password`
+
+### Example
+
+`example.py` workflow example:
+```python
+from flytekit import ImageSpec, task, workflow
+from flytekitplugins.perian_job import PerianConfig
+
+image_spec = ImageSpec(
+    name="flyte-test",
+    registry="my-registry",
+    python_version="3.11",
+    apt_packages=["wget", "curl", "git"],
+    packages=[
+        "flytekitplugins-perian-job",
+    ],
+)
+
+@task(container_image=image_spec,
+      task_config=PerianConfig(
+        accelerators=1,
+        accelerator_type="A100",
+      ))
+def perian_hello(name: str) -> str:
+    return f"hello {name}!"
+
+@workflow
+def my_wf(name: str = "world") -> str:
+    return perian_hello(name=name)
+```

flytekitplugins-perian_job-1.12.0/flytekitplugins/perian_job/__init__.py

@@ -0,0 +1,2 @@
+from .agent import PerianAgent
+from .task import PerianConfig, PerianTask

flytekitplugins-perian_job-1.12.0/flytekitplugins/perian_job/agent.py

@@ -0,0 +1,200 @@
+import shlex
+from dataclasses import dataclass
+from typing import Optional
+
+from flyteidl.core.execution_pb2 import TaskExecution
+from perian import (
+    AcceleratorQueryInput,
+    ApiClient,
+    Configuration,
+    CpuQueryInput,
+    CreateJobRequest,
+    DockerRegistryCredentials,
+    DockerRunParameters,
+    InstanceTypeQueryInput,
+    JobApi,
+    JobStatus,
+    MemoryQueryInput,
+    Name,
+    ProviderQueryInput,
+    RegionQueryInput,
+    Size,
+)
+
+from flytekit import current_context
+from flytekit.exceptions.base import FlyteException
+from flytekit.exceptions.user import FlyteUserException
+from flytekit.extend.backend.base_agent import AgentRegistry, AsyncAgentBase, Resource, ResourceMeta
+from flytekit.loggers import logger
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+PERIAN_API_URL = "https://api.perian.cloud"
+
+
+@dataclass
+class PerianMetadata(ResourceMeta):
+    """Metadata for Perian jobs"""
+
+    job_id: str
+
+
+class PerianAgent(AsyncAgentBase):
+    """Flyte Agent for executing tasks on Perian"""
+
+    name = "Perian Agent"
+
+    def __init__(self):
+        logger.info("Initializing Perian agent")
+        super().__init__(task_type_name="perian_task", metadata_type=PerianMetadata)
+
+    async def create(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap],
+        output_prefix: Optional[str],
+        **kwargs,
+    ) -> PerianMetadata:
+        logger.info("Creating new Perian job")
+        logger.debug("Task template: %s", task_template.__dict__)
+
+        config = Configuration(host=PERIAN_API_URL)
+        job_request = self._build_create_job_request(task_template)
+        with ApiClient(config) as api_client:
+            api_instance = JobApi(api_client)
+            response = api_instance.create_job(
+                create_job_request=job_request,
+                _headers=self._build_headers(),
+            )
+        if response.status_code != 200:
+            raise FlyteException(f"Failed to create Perian job: {response.text}")
+
+        return PerianMetadata(job_id=response.id)
+
+    def get(self, resource_meta: PerianMetadata, **kwargs) -> Resource:
+        job_id = resource_meta.job_id
+        logger.info("Getting Perian job status: %s", job_id)
+        config = Configuration(host=PERIAN_API_URL)
+        with ApiClient(config) as api_client:
+            api_instance = JobApi(api_client)
+            response = api_instance.get_job_by_id(
+                job_id=str(job_id),
+                _headers=self._build_headers(),
+            )
+        if response.status_code != 200:
+            raise FlyteException(f"Failed to get Perian job status: {response.text}")
+        if not response.jobs:
+            raise FlyteException(f"Perian job not found: {job_id}")
+        job = response.jobs[0]
+
+        return Resource(
+            phase=self._perian_job_status_to_flyte_phase(job.status),
+            message=job.logs,
+        )
+
+    def delete(self, resource_meta: PerianMetadata, **kwargs):
+        job_id = resource_meta.job_id
+        logger.info("Cancelling Perian job: %s", job_id)
+        config = Configuration(host=PERIAN_API_URL)
+        with ApiClient(config) as api_client:
+            api_instance = JobApi(api_client)
+            response = api_instance.cancel_job(
+                job_id=str(job_id),
+                _headers=self._build_headers(),
+            )
+        if response.status_code != 200:
+            raise FlyteException(f"Failed to cancel Perian job: {response.text}")
+
+    def _build_create_job_request(self, task_template: TaskTemplate) -> CreateJobRequest:
+        params = task_template.custom
+        secrets = current_context().secrets
+
+        # Build instance type requirements
+        reqs = InstanceTypeQueryInput()
+        if params.get("cores"):
+            reqs.cpu = CpuQueryInput(cores=int(params["cores"]))
+        if params.get("memory"):
+            reqs.ram = MemoryQueryInput(size=Size(params["memory"]))
+        if any([params.get("accelerators"), params.get("accelerator_type")]):
+            reqs.accelerator = AcceleratorQueryInput()
+            if params.get("accelerators"):
+                reqs.accelerator.no = int(params["accelerators"])
+            if params.get("accelerator_type"):
+                reqs.accelerator.name = Name(params["accelerator_type"])
+        if params.get("country_code"):
+            reqs.region = RegionQueryInput(location=params["country_code"])
+        if params.get("provider"):
+            reqs.provider = ProviderQueryInput(name_short=params["provider"])
+
+        docker_run = DockerRunParameters()
+        # Credentials to the Flyte storage S3 bucket need to be passed to the Perian job.
+        aws_access_key_id = secrets.get("aws_access_key_id")
+        aws_secret_access_key = secrets.get("aws_secret_access_key")
+        if not aws_access_key_id or not aws_secret_access_key:
+            raise FlyteUserException(
+                "AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY for the Flyte storage bucket "
+                "must be provided in the secrets"
+            )
+        docker_run.env_variables = {
+            "AWS_ACCESS_KEY_ID": aws_access_key_id,
+            "AWS_SECRET_ACCESS_KEY": aws_secret_access_key,
+        }
+
+        docker_registry = None
+        try:
+            dr_url = secrets.get("docker_registry_url")
+            dr_username = secrets.get("docker_registry_username")
+            dr_password = secrets.get("docker_registry_password")
+            if any([dr_url, dr_username, dr_password]):
+                docker_registry = DockerRegistryCredentials(
+                    url=dr_url,
+                    username=dr_username,
+                    password=dr_password,
+                )
+        except ValueError:
+            pass
+
+        container = task_template.container
+        if ":" in container.image:
+            docker_run.image_name, docker_run.image_tag = container.image.rsplit(":", 1)
+        else:
+            docker_run.image_name = container.image
+        if container.args:
+            docker_run.command = shlex.join(container.args)
+
+        return CreateJobRequest(
+            auto_failover_instance_type=True,
+            requirements=reqs,
+            docker_run_parameters=docker_run,
+            docker_registry_credentials=docker_registry,
+        )
+
+    def _build_headers(self) -> dict:
+        secrets = current_context().secrets
+        org = secrets.get("perian_organization")
+        token = secrets.get("perian_token")
+        if not org or not token:
+            raise FlyteUserException(
+                "perian_organization and perian_token must be provided in the secrets")
+        return {
+            "X-PERIAN-AUTH-ORG": org,
+            "Authorization": "Bearer " + token,
+        }
+
+    def _perian_job_status_to_flyte_phase(self, status: JobStatus) -> TaskExecution.Phase:
+        status_map = {
+            JobStatus.QUEUED: TaskExecution.QUEUED,
+            JobStatus.INITIALIZING: TaskExecution.INITIALIZING,
+            JobStatus.RUNNING: TaskExecution.RUNNING,
+            JobStatus.DONE: TaskExecution.SUCCEEDED,
+            JobStatus.SERVERERROR: TaskExecution.FAILED,
+            JobStatus.USERERROR: TaskExecution.FAILED,
+            JobStatus.CANCELLED: TaskExecution.ABORTED,
+        }
+        if status == JobStatus.UNDEFINED:
+            raise FlyteException("Undefined Perian job status")
+        return status_map[status]
+
+
+# To register the Perian agent
+AgentRegistry.register(PerianAgent())

flytekitplugins-perian_job-1.12.0/flytekitplugins/perian_job/task.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional, Union
+
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+
+from flytekit import FlyteContextManager, PythonFunctionTask, logger
+from flytekit.configuration import SerializationSettings
+from flytekit.exceptions.user import FlyteUserException
+from flytekit.extend import TaskPlugins
+from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
+from flytekit.image_spec import ImageSpec
+
+
+@dataclass
+class PerianConfig:
+    """Used to configure a Perian Task"""
+
+    # Number of CPU cores
+    cores: Optional[int] = None
+    # Amount of memory in GB
+    memory: Optional[int] = None
+    # Number of accelerators
+    accelerators: Optional[int] = None
+    # Type of accelerator (e.g. 'A100')
+    # For a full list of supported accelerators, use the perian CLI list-accelerators command
+    accelerator_type: Optional[str] = None
+    # Country code to run the job in (e.g. 'DE')
+    country_code: Optional[str] = None
+    # Cloud provider to run the job on
+    provider: Optional[str] = None
+
+
+class PerianTask(AsyncAgentExecutorMixin, PythonFunctionTask):
+    """A special task type for running tasks on Perian Job Platform (perian.io)"""
+
+    _TASK_TYPE = "perian_task"
+
+    def __init__(
+        self,
+        task_config: Optional[PerianConfig],
+        task_function: Callable,
+        container_image: Optional[Union[str, ImageSpec]] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            task_config=task_config,
+            task_function=task_function,
+            container_image=container_image,
+            task_type=self._TASK_TYPE,
+            **kwargs,
+        )
+
+    def execute(self, **kwargs) -> Any:
+        if isinstance(self.task_config, PerianConfig):
+            # Use the Perian agent to run it by default.
+            try:
+                ctx = FlyteContextManager.current_context()
+                if not ctx.file_access.is_remote(ctx.file_access.raw_output_prefix):
+                    raise ValueError(
+                        "To submit a Perian job locally,"
+                        " please set --raw-output-data-prefix to a remote path. e.g. s3://, gcs//, etc."
+                    )
+                if ctx.execution_state and ctx.execution_state.is_local_execution():
+                    return AsyncAgentExecutorMixin.execute(self, **kwargs)
+            except Exception as e:
+                logger.error("Agent failed to run the task with error: %s", e)
+                raise
+        return PythonFunctionTask.execute(self, **kwargs)
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        """
+        Return plugin-specific data as a serializable dictionary.
+        """
+        config = {
+            "cores": self.task_config.cores,
+            "memory": self.task_config.memory,
+            "accelerators": self.task_config.accelerators,
+            "accelerator_type": self.task_config.accelerator_type,
+            "country_code": _validate_and_format_country_code(self.task_config.country_code),
+            "provider": self.task_config.provider,
+        }
+        config = {k: v for k, v in config.items() if v is not None}
+        s = Struct()
+        s.update(config)
+        return json_format.MessageToDict(s)
+
+
+def _validate_and_format_country_code(country_code: Optional[str]) -> Optional[str]:
+    if not country_code:
+        return None
+    if len(country_code) != 2:
+        raise FlyteUserException("Invalid country code. Please provide a valid two-letter country code. (e.g. DE)")
+    return country_code.upper()
+
+
+# Inject the Perian plugin into flytekits dynamic plugin loading system
+TaskPlugins.register_pythontask_plugin(PerianConfig, PerianTask)

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/PKG-INFO

@@ -0,0 +1,92 @@
+Metadata-Version: 2.1
+Name: flytekitplugins-perian-job
+Version: 1.12.0
+Summary: Flyte agent for Perian Job Platform (perian.io)
+Author: Omar Tarabai
+Author-email: otarabai@perian.io
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+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: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+
+# Flytekit Perian Job Platform Plugin
+
+Flyte Agent plugin for executing Flyte tasks on Perian Job Platform (perian.io).
+
+Perian Job Platform is still in closed beta. Contact support@perian.io if you are interested in trying it out.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-perian-job
+```
+
+## Getting Started
+
+This plugin allows executing `PythonFunctionTask` on Perian.
+
+An [ImageSpec](https://docs.flyte.org/en/latest/user_guide/customizing_dependencies/imagespec.html) need to be built with the perian agent plugin installed.
+
+### Parameters
+
+The following parameters can be used to set the requirements for the Perian task. If any of the requirements are skipped, it is replaced with the cheapest option. At least one requirement value should be set.
+* `cores`: Number of CPU cores
+* `memory`: Amount of memory in GB
+* `accelerators`: Number of accelerators
+* `accelerator_type`: Type of accelerator (e.g. 'A100'). For a full list of supported accelerators, use the perian CLI list-accelerators command.
+* `country_code`: Country code to run the job in (e.g. 'DE')
+
+### Credentials
+
+The following [secrets](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html) are required to be defined for the agent server:
+* Perian credentials:
+    * `perian_organization`
+    * `perian_token`
+* AWS credentials for accessing the Flyte storage bucket. You might need to create long-lived credentials to avoid expiry when running tasks. These credentials are never logged by Perian and are stored only until it is used, then immediately deleted:
+    * `aws_access_key_id`
+    * `aws_secret_access_key`
+* (Optional) Custom docker registry for pulling the Flyte image:
+    * `docker_registry_url`
+    * `docker_registry_username`
+    * `docker_registry_password`
+
+### Example
+
+`example.py` workflow example:
+```python
+from flytekit import ImageSpec, task, workflow
+from flytekitplugins.perian_job import PerianConfig
+
+image_spec = ImageSpec(
+    name="flyte-test",
+    registry="my-registry",
+    python_version="3.11",
+    apt_packages=["wget", "curl", "git"],
+    packages=[
+        "flytekitplugins-perian-job",
+    ],
+)
+
+@task(container_image=image_spec,
+      task_config=PerianConfig(
+        accelerators=1,
+        accelerator_type="A100",
+      ))
+def perian_hello(name: str) -> str:
+    return f"hello {name}!"
+
+@workflow
+def my_wf(name: str = "world") -> str:
+    return perian_hello(name=name)
+```

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+setup.py
+flytekitplugins/perian_job/__init__.py
+flytekitplugins/perian_job/agent.py
+flytekitplugins/perian_job/task.py
+flytekitplugins_perian_job.egg-info/PKG-INFO
+flytekitplugins_perian_job.egg-info/SOURCES.txt
+flytekitplugins_perian_job.egg-info/dependency_links.txt
+flytekitplugins_perian_job.egg-info/entry_points.txt
+flytekitplugins_perian_job.egg-info/namespace_packages.txt
+flytekitplugins_perian_job.egg-info/requires.txt
+flytekitplugins_perian_job.egg-info/top_level.txt
+tests/__init__.py
+tests/test_perian.py

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flytekit.plugins]
+perian_job = flytekitplugins.perian_job

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/namespace_packages.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/requires.txt

@@ -0,0 +1,2 @@
+flytekit<2.0.0,>=1.12.0
+perian>=0.2.3

flytekitplugins-perian_job-1.12.0/flytekitplugins_perian_job.egg-info/top_level.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins-perian_job-1.12.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flytekitplugins-perian_job-1.12.0/setup.py

@@ -0,0 +1,39 @@
+from setuptools import setup
+
+PLUGIN_NAME = "perian_job"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>=1.12.0,<2.0.0", "perian>=0.2.3"]
+
+__version__ = "1.12.0"
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="Omar Tarabai",
+    author_email="otarabai@perian.io",
+    description="Flyte agent for Perian Job Platform (perian.io)",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.8",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

flytekitplugins-perian_job-1.12.0/tests/test_perian.py

@@ -0,0 +1,43 @@
+from collections import OrderedDict
+
+from flytekitplugins.perian_job import PerianConfig, PerianTask
+
+from flytekit import task
+from flytekit.configuration import DefaultImages, ImageConfig, SerializationSettings
+from flytekit.extend import get_serializable
+
+
+def test_perian_task():
+    task_config = PerianConfig(
+        cores=2,
+        memory="8",
+        accelerators=1,
+        accelerator_type="A100",
+        country_code="DE",
+    )
+    container_image = DefaultImages.default_image()
+
+    @task(
+        task_config=task_config,
+        container_image=container_image,
+    )
+    def say_hello(name: str) -> str:
+        return f"Hello, {name}."
+
+    assert say_hello.task_config == task_config
+    assert say_hello.task_type == "perian_task"
+    assert isinstance(say_hello, PerianTask)
+
+    serialization_settings = SerializationSettings(image_config=ImageConfig())
+    task_spec = get_serializable(OrderedDict(), serialization_settings, say_hello)
+    template = task_spec.template
+    container = template.container
+
+    assert template.custom == {
+        'accelerator_type': 'A100',
+        'accelerators': 1.0,
+        'cores': 2.0,
+        'country_code': 'DE',
+        'memory': '8',
+    }
+    assert container.image == container_image