Add mout section to allow specifying uploads via polyaxonfile

 * Allow uploading/mounting multiple folders via cli

Author: polyaxon-ci

cli/polyaxon/_cli/run.py

@@ -28,7 +28,7 @@
 from polyaxon._flow import V1Operation, V1RunPending
 from polyaxon._managers.git import GitConfigManager
 from polyaxon._managers.run import RunConfigManager
-from polyaxon._polyaxonfile import check_polyaxonfile
+from polyaxon._polyaxonfile import check_polyaxonfile, CompiledOperationSpecification
 from polyaxon._runner.kinds import RunnerKind
 from polyaxon._schemas.lifecycle import ManagedBy
 from polyaxon._utils import cache
@@ -161,14 +161,31 @@ def start_run_shell(run_uuid: str):
         }
         ctx.invoke(run_shell)
 
-    def upload_run(run_uuid: str):
+    def upload_run(run_uuid: str, path_from: str, path_to: str):
         ctx.obj = {
             "project": "{}/{}".format(owner_team, project_name),
             "run_uuid": run_uuid,
         }
         ctx.invoke(
-            run_upload, path_to=upload_to, path_from=upload_from, sync_failure=True
+            run_upload,
+            path_to=path_to,
+            path_from=path_from,
+            sync_failure=True,
         )
+
+    def upload_mounts(run_uuid: str):
+        if upload:
+            upload_run(
+                build_uuid or run_instance.uuid,
+                path_to=upload_to,
+                path_from=upload_from,
+            )
+
+        compiled_spec = CompiledOperationSpecification.read(run_instance.content)
+        mounts = compiled_spec.get_resolved_mount()
+        for mount in mounts:
+            upload_run(run_uuid, path_from=mount.path_from, path_to=mount.path_to)
+
         if approve_after and approve_after > 0:
             Printer.print(f"Waiting {approve_after}s to approve the run...")
             time.sleep(approve_after)
@@ -184,10 +201,12 @@ def upload_run(run_uuid: str):
         upload_to = upload_to or DEFAULT_UPLOADS_PATH
         run_meta_info = run_meta_info or {}
         run_meta_info[META_UPLOAD_ARTIFACTS] = upload_to
+
+    is_uploading = upload or op_spec.has_mount()
     run_instance = create_run(
         managed_by=managed_by,
         meta_info=run_meta_info,
-        pending=V1RunPending.UPLOAD if upload else None,
+        pending=V1RunPending.UPLOAD if is_uploading else None,
     )
     if not run_instance:
         return
@@ -200,8 +219,8 @@ def upload_run(run_uuid: str):
         build_name = run_instance.settings.build.get("name")
         runs_to_watch.insert(0, RunWatchSpec(build_uuid, build_name))
 
-    if upload:
-        upload_run(build_uuid or run_instance.uuid)
+    if is_uploading:
+        upload_mounts(build_uuid or run_instance.uuid)
 
     if local:
         for instance in runs_to_watch:
@@ -277,8 +296,20 @@ def upload_run(run_uuid: str):
     "-u",
     is_flag=True,
     default=False,
-    help="To upload the working dir to run's artifacts path "
-    "as an init context before scheduling the run.",
+    help="Upload the current working directory to run's artifacts path "
+    "as an init context before scheduling the run. "
+    "Uploads to 'uploads/' by default.",
+)
+@click.option(
+    "--mount",
+    "-m",
+    "mounts",
+    multiple=True,
+    type=str,
+    help="Mount local paths to run artifacts. "
+    "Syntax: './src:dest' or './src' (defaults to 'uploads/' destination). "
+    "Can be specified multiple times. "
+    "Examples: -m ./code -m ./data:datasets -m ./models:models",
 )
 @click.option(
     "--upload-from",
@@ -437,6 +468,7 @@ def run(
     shell,
     log,
     upload,
+    mounts,
     upload_from,
     upload_to,
     watch,
@@ -498,21 +530,20 @@ def run(
     $ polyaxon run -pm path/to/my-component.py:componentA
 
 
-    Uploading from everything in the current folder to the default uploads path
+    Uploading current folder to the default uploads path
 
     \b
     $ polyaxon run ... -u
 
-
-    Uploading from everything in the current folder to a custom path, e.g. code
+    Uploading from a subfolder to a custom path
 
     \b
-    $ polyaxon run ... -u-to code
+    $ polyaxon run ... -u -u-from ./code -u-to code
 
-    Uploading from everything from a sub-folder, e.g. ./code to the a custom path, e.g. new-code
+    Mounting multiple paths using -m
 
     \b
-    $ polyaxon run ... -u-from ./code -u-to new-code
+    $ polyaxon run ... -m ./src:code -m ./data:datasets -m ./models
     """
     if log and shell:
         Printer.error(
@@ -590,6 +621,13 @@ def run(
     owner, team, project_name = get_project_or_local(project, is_cli=True)
     tags = validate_tags(tags, validate_yaml=True)
 
+    # Handle CLI uploads - merge with polyaxonfile mount section
+    cli_mounts = list(mounts) if mounts else []
+
+    if cli_mounts:
+        existing_mount = list(op_spec.mount) if op_spec.mount else []
+        op_spec.mount = existing_mount + cli_mounts
+
     _run(
         ctx=ctx,
         name=name,

cli/polyaxon/_flow/component/base.py

@@ -14,6 +14,7 @@
 from polyaxon._flow.builds import V1Build
 from polyaxon._flow.cache import V1Cache
 from polyaxon._flow.hooks import V1Hook
+from polyaxon._schemas.types.mounts import V1Mount
 from polyaxon._flow.plugins import V1Plugins
 from polyaxon._flow.termination import V1Termination
 from polyaxon._schemas.base import BaseSchemaModel
@@ -35,9 +36,25 @@ class BaseComponent(BaseSchemaModel):
     hooks: Optional[Union[List[V1Hook], RefField]] = None
     is_approved: Optional[BoolOrRef] = Field(alias="isApproved", default=None)
     cost: Optional[FloatOrRef] = None
+    mount: Optional[List[Union[StrictStr, V1Mount]]] = None
 
     @field_validator("tags", "presets", **validation_before)
     def validate_str_list(cls, v):
         if isinstance(v, str):
             return to_list(v, check_str=True)
         return v
+
+    def get_resolved_mount(self):
+        if not self.mount:
+            return []
+        resolved_mounts = []
+        for mount in self.mount:
+            if isinstance(mount, str):
+                if ":" in mount:
+                    p_from, p_to = mount.split(":", 1)
+                else:
+                    p_from, p_to = mount, None
+                resolved_mounts.append(V1Mount(path_from=p_from, path_to=p_to))
+            else:
+                resolved_mounts.append(mount)
+        return resolved_mounts

cli/polyaxon/_flow/component/component.py

@@ -43,6 +43,7 @@ class V1Component(
         cache: [V1Cache](/docs/automation/helpers/cache/), optional
         termination: [V1Termination](/docs/core/specification/termination/), optional
         plugins: [V1Plugins](/docs/core/specification/plugins/), optional
+        mount: List[[V1Mount](/docs/core/specification/mount/)], optional
         build: [V1Build](/docs/automation/builds/), optional
         hooks: List[[V1Hook](/docs/automation/hooks/)], optional
         inputs: [V1IO](/docs/core/specification/io/), optional
@@ -69,6 +70,7 @@ class V1Component(
     >>>   hooks:
     >>>   inputs:
     >>>   outputs:
+    >>>   mount:
     >>>   build:
     >>>   run:
     >>>   isApproved:
@@ -94,6 +96,7 @@ class V1Component(
     >>>     hooks=[V1Hook(...)],
     >>>     inputs=[V1IO(...)],
     >>>     outputs=[V1IO(...)],
+    >>>     mount=[V1Mount(...)],
     >>>     build=V1Build(...),
     >>>     run=...
     >>> )
@@ -317,6 +320,22 @@ class V1Component(
     >>>   ...
     ```
 
+    ### mount
+
+    > **Note**: ver 2.13+. Please check [V1Mount](/docs/core/specification/mount/) for more details.
+
+    This section defines a list of mounts to be used for this operation.
+    Mounts can be defined either as strings or as full objects.
+    ```yaml
+    >>> operation:
+    >>>   ...
+    >>>   mount:
+    >>>     - /path/in/host:/path/in/container  # defined as string
+    >>>     - path_from: /path/in/host          # defined as object
+    >>>       path_to: /path/in/container
+    >>>   ...
+    ```
+
     ### build
 
     > **Note**: Please check [V1Build](/docs/automation/builds/) for more details.

cli/polyaxon/_flow/operations/base.py

@@ -14,7 +14,7 @@
 
 
 class BaseOp(BaseComponent, MatrixMixin, ScheduleMixin):
-    _FIELDS_SAME_KIND_PATCH = ["schedule", "matrix"]
+    _FIELDS_SAME_KIND_PATCH = ["schedule", "matrix", "mount"]
 
     schedule: Optional[V1Schedule] = None
     events: Optional[Union[List[V1EventTrigger], RefField]] = None

cli/polyaxon/_flow/operations/operation.py

@@ -58,6 +58,7 @@ class V1Operation(BaseOp, TemplateMixinConfig):
         params: Dict[str, [V1Param](/docs/core/specification/params/)], optional
         schedule: Union[[V1CronSchedule](/docs/automation/schedules/cron/), [V1IntervalSchedule](/docs/automation/schedules/interval/), [V1DateTimeSchedule](/docs/automation/schedules/datetime/)], optional  # noqa
         events: List[[V1EventTrigger](/docs/automation/events/)], optional
+        mount: List[[V1Mount](/docs/core/specification/mount/)], optional
         build: [V1Build](/docs/automation/builds/), optional
         hooks: List[[V1Hook](/docs/automation/hooks/)], optional
         matrix: Union[[V1Mapping](/docs/automation/mapping/), [V1GridSearch](/docs/automation/optimization-engine/grid-search/), [V1RandomSearch](/docs/automation/optimization-engine/random-search/), [V1Hyperband](/docs/automation/optimization-engine/hyperband/), [V1Bayes](/docs/automation/optimization-engine/bayesian-optimization/), [V1Hyperopt](/docs/automation/optimization-engine/hyperopt/), [V1Iterative](/docs/automation/optimization-engine/iterative/)], optional  # noqa
@@ -96,6 +97,7 @@ class V1Operation(BaseOp, TemplateMixinConfig):
     >>>   actions:
     >>>   hooks:
     >>>   params:
+    >>>   mount:
     >>>   build:
     >>>   runPatch:
     >>>   hubRef:
@@ -126,6 +128,7 @@ class V1Operation(BaseOp, TemplateMixinConfig):
     >>>     events=["event-ref1", "event-ref2"],
     >>>     hooks=[V1Hook(...)],
     >>>     outputs={"param1": V1Param(...), ...},
+    >>>     mount=[V1Mount(...)],
     >>>     build=V1Build(...),
     >>>     component=V1Component(...),
     >>> )
@@ -328,6 +331,22 @@ class V1Operation(BaseOp, TemplateMixinConfig):
     >>>   ...
     ```
 
+    ### mount
+
+    > **Note**: ver 2.13+. Please check [V1Mount](/docs/core/specification/mount/) for more details.
+
+    This section defines a list of mounts to be used for this operation.
+    Mounts can be defined either as strings or as full objects.
+    ```yaml
+    >>> operation:
+    >>>   ...
+    >>>   mount:
+    >>>     - /path/in/host:/path/in/container  # defined as string
+    >>>     - path_from: /path/in/host          # defined as object
+    >>>       path_to: /path/in/container
+    >>>   ...
+    ```
+
     ### build
 
     > **Note**: Please check [V1Build](/docs/automation/builds/) for more details.
@@ -712,5 +731,12 @@ def from_build(cls, build: V1Build, contexts: Optional[Dict] = None):
 
         return cls.model_construct(**content)
 
+    def has_mount(self):
+        if self.mount:
+            return True
+        if self.component and self.component.mount:
+            return True
+        return False
+
 
 PartialV1Operation = to_partial(V1Operation)

cli/polyaxon/_polyaxonfile/specs/compiled_operation.py

@@ -332,6 +332,7 @@ def apply_preset(
             "events",
             "plugins",
             "termination",
+            "mount",
             "matrix",
             "joins",
             "schedule",

cli/polyaxon/_polyaxonfile/specs/operation.py

@@ -110,6 +110,7 @@ def get_context_io(c_name: str, c_io: V1Param, is_list=None):
             "trigger",
             "conditions",
             "skip_on_upstream_skip",
+            "mount",
         }
         patch_keys = patch_keys.intersection(config.model_fields_set)
         patch_data = {k: getattr(config, k) for k in patch_keys}

cli/polyaxon/_polyaxonfile/specs/sections.py

@@ -39,6 +39,7 @@ class Sections:
     PATH_REF = "pathRef"
     URL_REF = "urlRef"
     COMPONENT = "component"
+    MOUNT = "mount"
 
     SECTIONS = (
         VERSION,
@@ -80,6 +81,7 @@ class Sections:
         CONTEXTS,
         RUN,
         RUN_PATCH,
+        MOUNT,
     )
 
     PARSING_SECTIONS = (
@@ -97,6 +99,7 @@ class Sections:
         DEPENDENCIES,
         TRIGGER,
         CONDITIONS,
+        MOUNT,
         SKIP_ON_UPSTREAM_SKIP,
         PATCH_STRATEGY,
         "is_approved",

cli/polyaxon/_schemas/types/mounts.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+from clipped.compact.pydantic import Field, StrictStr
+
+from polyaxon._schemas.types.base import BaseTypeConfig
+
+
+class V1Mount(BaseTypeConfig):
+    _IDENTIFIER = "mount"
+
+    path_from: Optional[StrictStr] = Field(alias="from", default=None)
+    path_to: Optional[StrictStr] = Field(alias="to", default=None)

cli/tests/test_polyaxonfile/test_mounts.py

@@ -0,0 +1,51 @@
+import pytest
+from polyaxon._flow import V1Operation, V1RunKind
+from polyaxon._schemas.types.mounts import V1Mount
+from polyaxon._utils.test_utils import BaseTestCase
+
+
+@pytest.mark.polyaxonfile_mark
+class TestMounts(BaseTestCase):
+    def test_mount_short_syntax(self):
+        config_dict = {
+            "version": 1.1,
+            "kind": "operation",
+            "mount": ["./src:/dst", "./src2"],
+            "component": {
+                "run": {"kind": V1RunKind.JOB, "container": {"image": "test"}}
+            },
+        }
+        op_config = V1Operation.read(values=config_dict)
+        assert op_config.mount == ["./src:/dst", "./src2"]
+        assert op_config.get_resolved_mount() == [
+            V1Mount(path_from="./src", path_to="/dst"),
+            V1Mount(path_from="./src2", path_to=None),
+        ]
+
+    def test_mount_long_syntax(self):
+        config_dict = {
+            "version": 1.1,
+            "kind": "operation",
+            "mount": [
+                {"from": "./src", "to": "/dst"},
+                {"from": "./src2"},
+                {"to": "/dst2"},
+            ],
+            "component": {
+                "run": {"kind": V1RunKind.JOB, "container": {"image": "test"}}
+            },
+        }
+        op_config = V1Operation.read(values=config_dict)
+        assert len(op_config.mount) == 3
+        assert isinstance(op_config.mount[0], V1Mount)
+        assert op_config.mount[0].path_from == "./src"
+        assert op_config.mount[0].path_to == "/dst"
+        assert op_config.mount[1].path_from == "./src2"
+        assert op_config.mount[1].path_to is None
+        assert op_config.mount[2].path_from is None
+        assert op_config.mount[2].path_to == "/dst2"
+        assert op_config.get_resolved_mount() == [
+            V1Mount(path_from="./src", path_to="/dst"),
+            V1Mount(path_from="./src2", path_to=None),
+            V1Mount(path_from=None, path_to="/dst2"),
+        ]