nodekit 0.2.0__py3-none-any.whl

nodekit/experimental/s3.py

@@ -0,0 +1,219 @@
+import pydantic
+
+import boto3
+import botocore.client
+import botocore.exceptions
+from typing import BinaryIO
+import hashlib
+import mimetypes
+
+from nodekit._internal.types.value import SHA256, MediaType
+import os
+from pathlib import Path
+
+from urllib.parse import quote
+
+
+# %%
+class UploadAssetResult(pydantic.BaseModel):
+    sha256: SHA256
+    mime_type: MediaType
+    asset_url: pydantic.HttpUrl
+
+
+class S3Client:
+    class Config(pydantic.BaseModel):
+        bucket_name: str
+        region_name: str
+        aws_access_key_id: str
+        aws_secret_access_key: pydantic.SecretStr
+
+    def __init__(
+        self,
+        bucket_name: str,
+        region_name: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+    ):
+        self.config = self.Config(
+            bucket_name=bucket_name,
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=pydantic.SecretStr(aws_secret_access_key),
+        )
+        self._client: botocore.client.BaseClient = boto3.client(
+            "s3",
+            region_name=self.config.region_name,
+            aws_access_key_id=self.config.aws_access_key_id,
+            aws_secret_access_key=self.config.aws_secret_access_key.get_secret_value(),
+        )
+
+    @staticmethod
+    def _derive_s3_key(
+        sha256: SHA256,
+        mime_type: MediaType,
+    ) -> str:
+        ext = mimetypes.guess_extension(mime_type)
+        if ext is None:
+            raise ValueError(
+                f"Could not determine file extension for mime type {mime_type}"
+            )
+        return f"assets/{mime_type}/{sha256}{ext}"
+
+    def _assemble_s3_url(self, key: str) -> pydantic.HttpUrl:
+        url = f"https://{self.config.bucket_name}.s3.{self.config.region_name}.amazonaws.com/{key}"
+        return pydantic.HttpUrl(url)
+
+    def maybe_resolve_asset(
+        self,
+        sha256: SHA256,
+        mime_type: MediaType,
+    ) -> UploadAssetResult | None:
+        # Derive S3 key
+        key = self._derive_s3_key(sha256=sha256, mime_type=mime_type)
+
+        # Check if it exists
+        try:
+            self._client.head_object(Bucket=self.config.bucket_name, Key=key)
+        except botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                return None
+            else:
+                raise RuntimeError(
+                    f"S3 head_object failed: {e.response.get('Error', {}).get('Message', str(e))}"
+                ) from e
+
+        # Return resolved asset ref
+        return UploadAssetResult(
+            mime_type=mime_type,
+            sha256=sha256,
+            asset_url=self._assemble_s3_url(key),
+        )
+
+    def upload_asset(
+        self,
+        claimed_sha256: SHA256,
+        file: BinaryIO,
+        mime_type: MediaType,
+    ) -> UploadAssetResult:
+        """
+        Stream `file` directly to S3 with the given MIME type, and returns it public link.
+        """
+        # ---- derive a safe S3 key
+        ext = mimetypes.guess_extension(mime_type)
+        if ext is None:
+            raise ValueError(
+                f"Could not determine file extension for mime type {mime_type}"
+            )
+
+        # First I/O pass: compute SHA256
+        try:
+            file.seek(0)
+        except Exception:
+            pass
+
+        sha256 = hashlib.sha256()
+        try:
+            file.seek(0)
+        except Exception:
+            pass
+        while True:
+            chunk = file.read(8192)  # 8 MB
+            if not chunk:
+                break
+            sha256.update(chunk)
+        sha256 = sha256.hexdigest()
+
+        if not claimed_sha256 == sha256:
+            raise ValueError(
+                f"SHA256 mismatch: claimed {claimed_sha256}, computed {sha256}"
+            )
+
+        # Return immediately if already uploaded:
+        key = self._derive_s3_key(sha256=sha256, mime_type=mime_type)
+        asset_url = self._assemble_s3_url(key)
+
+        # Second I/O pass: upload to S3
+        file.seek(0)
+
+        try:
+            self._client.upload_fileobj(
+                Fileobj=file,
+                Bucket=self.config.bucket_name,
+                Key=key,
+                ExtraArgs={
+                    "ContentType": str(mime_type),
+                    "ACL": "public-read",
+                },
+            )
+        except botocore.exceptions.ClientError as e:
+            raise RuntimeError(
+                f"S3 upload failed: {e.response.get('Error', {}).get('Message', str(e))}"
+            ) from e
+
+        # Package return model
+        return UploadAssetResult(
+            sha256=sha256,
+            asset_url=asset_url,
+            mime_type=mime_type,
+        )
+
+    def sync_file(
+        self,
+        local_path: os.PathLike | str,
+        local_root: os.PathLike | str,
+        bucket_root: os.PathLike | str,
+        force: bool = False,
+    ) -> str:
+        """
+        Upload a single file to S3 under the key derived from its path relative to `local_root`,
+        prefixed by `bucket_root`. Skips upload if an object already exists with the same size,
+        unless `force=True`. Makes the object public-read and sets ContentType.
+
+        Returns the public URL of the object.
+        """
+        p = Path(local_path)
+        local_root = Path(local_root)
+        if not p.is_file() or p.is_symlink():
+            raise ValueError(f"Unsupported file type: {p}")
+
+        # Derive S3 key (prefix + rel path), then URL-encode it (keep '/' separators)
+        rel = p.relative_to(local_root).as_posix()
+        prefix = str(bucket_root).strip("/")
+        key = f"{prefix}/{rel}" if prefix else rel
+        encoded_key = quote(key, safe="/")
+
+        # Compute URL (prefer virtual-hosted-style for AWS, path-style for custom endpoints)
+        endpoint = self._client.meta.endpoint_url.rstrip("/")
+        if "amazonaws.com" in endpoint:
+            url = f"https://{self.config.bucket_name}.s3.{self.config.region_name}.amazonaws.com/{encoded_key}"
+        else:
+            # e.g., MinIO or custom S3-compatible endpoint
+            url = f"{endpoint}/{self.config.bucket_name}/{encoded_key}"
+
+        # Skip if already present with same size (unless force)
+        try:
+            head = self._client.head_object(Bucket=self.config.bucket_name, Key=key)
+            if head.get("ContentLength") == p.stat().st_size and not force:
+                return url
+        except self._client.exceptions.ClientError as e:
+            if e.response.get("Error", {}).get("Code") != "404":
+                raise  # unexpected error → bubble up
+
+        # Upload
+        mime, _ = mimetypes.guess_type(p.name)
+        extra = {
+            "ContentType": mime or "application/octet-stream",
+            "ACL": "public-read",
+        }
+        with p.open("rb") as f:
+            self._client.upload_fileobj(
+                Fileobj=f,
+                Bucket=self.config.bucket_name,
+                Key=key,
+                ExtraArgs=extra,
+            )
+
+        print(f"Uploaded {p.name} to S3")
+
+        return url

nodekit/experimental/turk_helper.py

@@ -0,0 +1,223 @@
+import os
+import uuid
+from decimal import Decimal
+from pathlib import Path
+from typing import Iterable
+import glob
+import pydantic
+
+import nodekit as nk
+from nodekit.experimental.recruitment_services.base import (
+    RecruiterServiceClient,
+    CreateHitRequest,
+    SendBonusPaymentRequest,
+)
+from nodekit.experimental.s3 import S3Client
+
+# %%
+type HitId = str
+type AssignmentId = str
+type WorkerId = str
+
+
+# %%
+class TraceResult(pydantic.BaseModel):
+    hit_id: HitId
+    assignment_id: AssignmentId
+    worker_id: WorkerId
+    trace: nk.Trace | None  # If None, validation failed
+
+
+class HitRequest(pydantic.BaseModel):
+    graph: nk.Graph
+    num_assignments: int
+    base_payment_usd: str
+    title: str
+    duration_sec: int = pydantic.Field(gt=0)
+    unique_request_token: str | None
+    hit_id: HitId
+
+
+class Helper:
+    """
+    Experimental; this might be moved to PsyHub / PsychoScope.
+    """
+
+    def __init__(
+        self,
+        recruiter_service_client: RecruiterServiceClient,
+        s3_client: S3Client,
+        local_cachedir: os.PathLike | str,
+    ):
+        self.recruiter_service_client = recruiter_service_client
+        self.s3_client = s3_client
+        self.local_cachedir = Path(local_cachedir)
+
+    def _get_hit_cachedir(self) -> Path:
+        return (
+            self.local_cachedir
+            / "hits"
+            / self.recruiter_service_client.get_recruiter_service_name()
+        )
+
+    def create_hit(
+        self,
+        graph: nk.Graph,
+        num_assignments: int,
+        base_payment_usd: str,
+        title: str,
+        duration_sec: int,
+        project_name: str,
+        unique_request_token: str | None = None,
+    ) -> HitId:
+        """
+        Creates a HIT based on the given Graph.
+        Automatically ensures a public site for the Graph exists on S3.
+        Caches the HIT (and its Graph) in the local cache.
+        """
+
+        graph_site_url = self.upload_graph_site(graph=graph)
+
+        if unique_request_token is None:
+            unique_request_token = uuid.uuid4().hex
+
+        response = self.recruiter_service_client.create_hit(
+            request=CreateHitRequest(
+                entrypoint_url=graph_site_url,
+                title=title,
+                description=title,
+                keywords=["psychology", "task", "cognitive", "science", "game"],
+                num_assignments=num_assignments,
+                duration_sec=duration_sec,
+                completion_reward_usd=Decimal(base_payment_usd),
+                unique_request_token=unique_request_token,
+                allowed_participant_ids=[],
+            )
+        )
+        hit_id: HitId = response.hit_id
+
+        # Just save the raw wire model, and hope the asset refs don't change. Todo: !
+        try:
+            hit_request = HitRequest(
+                graph=graph,
+                num_assignments=num_assignments,
+                base_payment_usd=base_payment_usd,
+                title=title,
+                duration_sec=duration_sec,
+                unique_request_token=unique_request_token,
+                hit_id=hit_id,
+            )
+            savepath = self._get_hit_cachedir() / project_name / f"{hit_id}.json"
+            if not savepath.parent.exists():
+                savepath.parent.mkdir(parents=True)
+            savepath.write_text(hit_request.model_dump_json(indent=2))
+        except Exception as e:
+            raise Exception(
+                f"Could not save Graph for HIT ({hit_id}) to local cache."
+            ) from e
+
+        return hit_id
+
+    def list_hits(self, project_name: str | None = None) -> list[HitId]:
+        # Just read off the local cache
+        savedir = self._get_hit_cachedir()
+        savedir.mkdir(parents=True, exist_ok=True)
+        hit_ids: list[HitId] = []
+
+        if project_name is None:
+            search_results = glob.glob(str(savedir / "**/*.json"), recursive=True)
+        else:
+            search_results = glob.glob(str(savedir / f"{project_name}/*.json"))
+
+        for path in search_results:
+            hit_ids.append(Path(path).stem.split("*.json")[0])
+        return hit_ids
+
+    def upload_graph_site(self, graph: nk.Graph) -> str:
+        """
+        Returns a URL to a public Graph site.
+        """
+
+        # Build the Graph site
+        build_site_result = nk.build_site(graph=graph, savedir=self.local_cachedir)
+
+        # Ensure index is sync'd
+        index_path = build_site_result.site_root / build_site_result.entrypoint
+        index_url = self.s3_client.sync_file(
+            local_path=index_path,
+            local_root=build_site_result.site_root,
+            bucket_root="",
+            force=False,
+        )
+
+        # Ensure deps are sync'd
+        for dep in build_site_result.dependencies:
+            self.s3_client.sync_file(
+                local_path=build_site_result.site_root / dep,
+                local_root=build_site_result.site_root,
+                bucket_root="",
+                force=False,
+            )
+
+        return index_url
+
+    def iter_traces(
+        self,
+        hit_id: HitId,
+    ) -> Iterable[TraceResult]:
+        """
+        Iterate the Traces collected under the given HIT ID.
+        Automatically approves any unapproved assignments.
+        """
+
+        # Pull new assignments
+        for asn in self.recruiter_service_client.iter_assignments(hit_id=hit_id):
+            # Ensure assignment is approved
+            if asn.status != "Approved":
+                self.recruiter_service_client.approve_assignment(
+                    assignment_id=asn.assignment_id,
+                )
+            try:
+                trace = nk.Trace.model_validate_json(asn.submission_payload)
+            except pydantic.ValidationError:
+                print(
+                    f"\n\n{asn.assignment_id}: Error validating submission payload:",
+                    asn.submission_payload,
+                )
+                trace = None
+
+            yield TraceResult(
+                hit_id=hit_id,
+                assignment_id=asn.assignment_id,
+                worker_id=asn.worker_id,
+                trace=trace,
+            )
+
+    def pay_bonus(
+        self,
+        worker_id: WorkerId,
+        assignment_id: AssignmentId,
+        amount_usd: str,
+    ) -> None:
+        self.recruiter_service_client.send_bonus_payment(
+            request=SendBonusPaymentRequest(
+                assignment_id=assignment_id,
+                amount_usd=Decimal(amount_usd),
+                worker_id=worker_id,
+            )
+        )
+
+    def get_hit(
+        self,
+        hit_id: HitId,
+    ) -> HitRequest:
+        """
+        Loads the Graph associated with the given HIT ID.
+        (Hit the local cache)
+        """
+        savepath = self._get_hit_cachedir() / f"{hit_id}.json"
+        if not savepath.parent.exists():
+            raise Exception(f"Could not save Graph for HIT {hit_id}.")
+
+        hit_request = HitRequest.model_validate_json(savepath.read_text())
+        return hit_request